#loc = loc(unknown)
module @pjit__wrapped_step_fn {
  func.func public @main(%arg0: tensor<ui32> {jax.arg_info = "state.step", mhlo.sharding = "{replicated}", tf.aliasing_output = 0 : i32} loc(unknown), %arg1: tensor<2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['final_ln']['bias']", mhlo.sharding = "{replicated}", tf.aliasing_output = 1 : i32} loc(unknown), %arg2: tensor<2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['final_ln']['scale']", mhlo.sharding = "{replicated}", tf.aliasing_output = 2 : i32} loc(unknown), %arg3: tensor<2048x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['position_emb']['emb_var']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 3 : i32} loc(unknown), %arg4: tensor<51200xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['softmax']['logits_ffn']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 4 : i32} loc(unknown), %arg5: tensor<2048x51200xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['softmax']['logits_ffn']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 5 : i32} loc(unknown), %arg6: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 6 : i32} loc(unknown), %arg7: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 7 : i32} loc(unknown), %arg8: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 8 : i32} loc(unknown), %arg9: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 9 : i32} loc(unknown), %arg10: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 10 : i32} loc(unknown), %arg11: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 11 : i32} loc(unknown), %arg12: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 12 : i32} loc(unknown), %arg13: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 13 : i32} loc(unknown), %arg14: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 14 : i32} loc(unknown), %arg15: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 15 : i32} loc(unknown), %arg16: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 16 : i32} loc(unknown), %arg17: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 17 : i32} loc(unknown), %arg18: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 18 : i32} loc(unknown), %arg19: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 19 : i32} loc(unknown), %arg20: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 20 : i32} loc(unknown), %arg21: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 21 : i32} loc(unknown), %arg22: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 22 : i32} loc(unknown), %arg23: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 23 : i32} loc(unknown), %arg24: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 24 : i32} loc(unknown), %arg25: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 25 : i32} loc(unknown), %arg26: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 26 : i32} loc(unknown), %arg27: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 27 : i32} loc(unknown), %arg28: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 28 : i32} loc(unknown), %arg29: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 29 : i32} loc(unknown), %arg30: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 30 : i32} loc(unknown), %arg31: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 31 : i32} loc(unknown), %arg32: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 32 : i32} loc(unknown), %arg33: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 33 : i32} loc(unknown), %arg34: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 34 : i32} loc(unknown), %arg35: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 35 : i32} loc(unknown), %arg36: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 36 : i32} loc(unknown), %arg37: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 37 : i32} loc(unknown), %arg38: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 38 : i32} loc(unknown), %arg39: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 39 : i32} loc(unknown), %arg40: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 40 : i32} loc(unknown), %arg41: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 41 : i32} loc(unknown), %arg42: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 42 : i32} loc(unknown), %arg43: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 43 : i32} loc(unknown), %arg44: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 44 : i32} loc(unknown), %arg45: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 45 : i32} loc(unknown), %arg46: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 46 : i32} loc(unknown), %arg47: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 47 : i32} loc(unknown), %arg48: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 48 : i32} loc(unknown), %arg49: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 49 : i32} loc(unknown), %arg50: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 50 : i32} loc(unknown), %arg51: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 51 : i32} loc(unknown), %arg52: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 52 : i32} loc(unknown), %arg53: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 53 : i32} loc(unknown), %arg54: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 54 : i32} loc(unknown), %arg55: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 55 : i32} loc(unknown), %arg56: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 56 : i32} loc(unknown), %arg57: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 57 : i32} loc(unknown), %arg58: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 58 : i32} loc(unknown), %arg59: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 59 : i32} loc(unknown), %arg60: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 60 : i32} loc(unknown), %arg61: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 61 : i32} loc(unknown), %arg62: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 62 : i32} loc(unknown), %arg63: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 63 : i32} loc(unknown), %arg64: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 64 : i32} loc(unknown), %arg65: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 65 : i32} loc(unknown), %arg66: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 66 : i32} loc(unknown), %arg67: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 67 : i32} loc(unknown), %arg68: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 68 : i32} loc(unknown), %arg69: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 69 : i32} loc(unknown), %arg70: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 70 : i32} loc(unknown), %arg71: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 71 : i32} loc(unknown), %arg72: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 72 : i32} loc(unknown), %arg73: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 73 : i32} loc(unknown), %arg74: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 74 : i32} loc(unknown), %arg75: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 75 : i32} loc(unknown), %arg76: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 76 : i32} loc(unknown), %arg77: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 77 : i32} loc(unknown), %arg78: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 78 : i32} loc(unknown), %arg79: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 79 : i32} loc(unknown), %arg80: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 80 : i32} loc(unknown), %arg81: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 81 : i32} loc(unknown), %arg82: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 82 : i32} loc(unknown), %arg83: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 83 : i32} loc(unknown), %arg84: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 84 : i32} loc(unknown), %arg85: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 85 : i32} loc(unknown), %arg86: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 86 : i32} loc(unknown), %arg87: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 87 : i32} loc(unknown), %arg88: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 88 : i32} loc(unknown), %arg89: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 89 : i32} loc(unknown), %arg90: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 90 : i32} loc(unknown), %arg91: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 91 : i32} loc(unknown), %arg92: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 92 : i32} loc(unknown), %arg93: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 93 : i32} loc(unknown), %arg94: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 94 : i32} loc(unknown), %arg95: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 95 : i32} loc(unknown), %arg96: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 96 : i32} loc(unknown), %arg97: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 97 : i32} loc(unknown), %arg98: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 98 : i32} loc(unknown), %arg99: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 99 : i32} loc(unknown), %arg100: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 100 : i32} loc(unknown), %arg101: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 101 : i32} loc(unknown), %arg102: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 102 : i32} loc(unknown), %arg103: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 103 : i32} loc(unknown), %arg104: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 104 : i32} loc(unknown), %arg105: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 105 : i32} loc(unknown), %arg106: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 106 : i32} loc(unknown), %arg107: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 107 : i32} loc(unknown), %arg108: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 108 : i32} loc(unknown), %arg109: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 109 : i32} loc(unknown), %arg110: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 110 : i32} loc(unknown), %arg111: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 111 : i32} loc(unknown), %arg112: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 112 : i32} loc(unknown), %arg113: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 113 : i32} loc(unknown), %arg114: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 114 : i32} loc(unknown), %arg115: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 115 : i32} loc(unknown), %arg116: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 116 : i32} loc(unknown), %arg117: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 117 : i32} loc(unknown), %arg118: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 118 : i32} loc(unknown), %arg119: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 119 : i32} loc(unknown), %arg120: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 120 : i32} loc(unknown), %arg121: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 121 : i32} loc(unknown), %arg122: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 122 : i32} loc(unknown), %arg123: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 123 : i32} loc(unknown), %arg124: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 124 : i32} loc(unknown), %arg125: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 125 : i32} loc(unknown), %arg126: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 126 : i32} loc(unknown), %arg127: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 127 : i32} loc(unknown), %arg128: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 128 : i32} loc(unknown), %arg129: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 129 : i32} loc(unknown), %arg130: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 130 : i32} loc(unknown), %arg131: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 131 : i32} loc(unknown), %arg132: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 132 : i32} loc(unknown), %arg133: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 133 : i32} loc(unknown), %arg134: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 134 : i32} loc(unknown), %arg135: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 135 : i32} loc(unknown), %arg136: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 136 : i32} loc(unknown), %arg137: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 137 : i32} loc(unknown), %arg138: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 138 : i32} loc(unknown), %arg139: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 139 : i32} loc(unknown), %arg140: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 140 : i32} loc(unknown), %arg141: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 141 : i32} loc(unknown), %arg142: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 142 : i32} loc(unknown), %arg143: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 143 : i32} loc(unknown), %arg144: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 144 : i32} loc(unknown), %arg145: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 145 : i32} loc(unknown), %arg146: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 146 : i32} loc(unknown), %arg147: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 147 : i32} loc(unknown), %arg148: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 148 : i32} loc(unknown), %arg149: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 149 : i32} loc(unknown), %arg150: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 150 : i32} loc(unknown), %arg151: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 151 : i32} loc(unknown), %arg152: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 152 : i32} loc(unknown), %arg153: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 153 : i32} loc(unknown), %arg154: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 154 : i32} loc(unknown), %arg155: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 155 : i32} loc(unknown), %arg156: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 156 : i32} loc(unknown), %arg157: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 157 : i32} loc(unknown), %arg158: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 158 : i32} loc(unknown), %arg159: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 159 : i32} loc(unknown), %arg160: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 160 : i32} loc(unknown), %arg161: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 161 : i32} loc(unknown), %arg162: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 162 : i32} loc(unknown), %arg163: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 163 : i32} loc(unknown), %arg164: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 164 : i32} loc(unknown), %arg165: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 165 : i32} loc(unknown), %arg166: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 166 : i32} loc(unknown), %arg167: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 167 : i32} loc(unknown), %arg168: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 168 : i32} loc(unknown), %arg169: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 169 : i32} loc(unknown), %arg170: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 170 : i32} loc(unknown), %arg171: tensor<1x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 171 : i32} loc(unknown), %arg172: tensor<1x2048x8192xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 172 : i32} loc(unknown), %arg173: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 173 : i32} loc(unknown), %arg174: tensor<1x8192x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 174 : i32} loc(unknown), %arg175: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 175 : i32} loc(unknown), %arg176: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 176 : i32} loc(unknown), %arg177: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 177 : i32} loc(unknown), %arg178: tensor<1x2048xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 178 : i32} loc(unknown), %arg179: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 179 : i32} loc(unknown), %arg180: tensor<1x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 180 : i32} loc(unknown), %arg181: tensor<1x2048x32x64xf32> {jax.arg_info = "state.mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 181 : i32} loc(unknown), %arg182: tensor<i32> {jax.arg_info = "state.opt_states[0][<flat index 0>][0][<flat index 0>]", mhlo.sharding = "{replicated}", tf.aliasing_output = 182 : i32} loc(unknown), %arg183: tensor<i32> {jax.arg_info = "state.opt_states[0][<flat index 0>][1][<flat index 0>]", mhlo.sharding = "{replicated}", tf.aliasing_output = 183 : i32} loc(unknown), %arg184: tensor<i32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 0>]", mhlo.sharding = "{replicated}", tf.aliasing_output = 184 : i32} loc(unknown), %arg185: tensor<2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['final_ln']['bias']", mhlo.sharding = "{replicated}", tf.aliasing_output = 185 : i32} loc(unknown), %arg186: tensor<2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['final_ln']['scale']", mhlo.sharding = "{replicated}", tf.aliasing_output = 186 : i32} loc(unknown), %arg187: tensor<2048x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['position_emb']['emb_var']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 187 : i32} loc(unknown), %arg188: tensor<51200xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['softmax']['logits_ffn']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 188 : i32} loc(unknown), %arg189: tensor<2048x51200xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['softmax']['logits_ffn']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 189 : i32} loc(unknown), %arg190: tensor<2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['final_ln']['bias']", mhlo.sharding = "{replicated}", tf.aliasing_output = 190 : i32} loc(unknown), %arg191: tensor<2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['final_ln']['scale']", mhlo.sharding = "{replicated}", tf.aliasing_output = 191 : i32} loc(unknown), %arg192: tensor<2048x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['position_emb']['emb_var']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 192 : i32} loc(unknown), %arg193: tensor<51200xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['softmax']['logits_ffn']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 193 : i32} loc(unknown), %arg194: tensor<2048x51200xf32> {jax.arg_info = "state.opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['softmax']['logits_ffn']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 194 : i32} loc(unknown), %arg195: tensor<i32> {jax.arg_info = "state.opt_states[0][<flat index 0>][3][<flat index 0>]", mhlo.sharding = "{replicated}", tf.aliasing_output = 195 : i32} loc(unknown), %arg196: tensor<1xi32> {jax.arg_info = "state.opt_states[0][<flat index 1>][0][<flat index 0>]", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 196 : i32} loc(unknown), %arg197: tensor<1xi32> {jax.arg_info = "state.opt_states[0][<flat index 1>][1][<flat index 0>]", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 197 : i32} loc(unknown), %arg198: tensor<1xi32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 0>]", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 198 : i32} loc(unknown), %arg199: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 199 : i32} loc(unknown), %arg200: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 200 : i32} loc(unknown), %arg201: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 201 : i32} loc(unknown), %arg202: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 202 : i32} loc(unknown), %arg203: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 203 : i32} loc(unknown), %arg204: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 204 : i32} loc(unknown), %arg205: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 205 : i32} loc(unknown), %arg206: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 206 : i32} loc(unknown), %arg207: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 207 : i32} loc(unknown), %arg208: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 208 : i32} loc(unknown), %arg209: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 209 : i32} loc(unknown), %arg210: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 210 : i32} loc(unknown), %arg211: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 211 : i32} loc(unknown), %arg212: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 212 : i32} loc(unknown), %arg213: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 213 : i32} loc(unknown), %arg214: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 214 : i32} loc(unknown), %arg215: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 215 : i32} loc(unknown), %arg216: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 216 : i32} loc(unknown), %arg217: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 217 : i32} loc(unknown), %arg218: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 218 : i32} loc(unknown), %arg219: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 219 : i32} loc(unknown), %arg220: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 220 : i32} loc(unknown), %arg221: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 221 : i32} loc(unknown), %arg222: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 222 : i32} loc(unknown), %arg223: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 223 : i32} loc(unknown), %arg224: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 224 : i32} loc(unknown), %arg225: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 225 : i32} loc(unknown), %arg226: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 226 : i32} loc(unknown), %arg227: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 227 : i32} loc(unknown), %arg228: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 228 : i32} loc(unknown), %arg229: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 229 : i32} loc(unknown), %arg230: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 230 : i32} loc(unknown), %arg231: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 231 : i32} loc(unknown), %arg232: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 232 : i32} loc(unknown), %arg233: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 233 : i32} loc(unknown), %arg234: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 234 : i32} loc(unknown), %arg235: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 235 : i32} loc(unknown), %arg236: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 236 : i32} loc(unknown), %arg237: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 237 : i32} loc(unknown), %arg238: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 238 : i32} loc(unknown), %arg239: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 239 : i32} loc(unknown), %arg240: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 240 : i32} loc(unknown), %arg241: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 241 : i32} loc(unknown), %arg242: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 242 : i32} loc(unknown), %arg243: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 243 : i32} loc(unknown), %arg244: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 244 : i32} loc(unknown), %arg245: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 245 : i32} loc(unknown), %arg246: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 246 : i32} loc(unknown), %arg247: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 247 : i32} loc(unknown), %arg248: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 248 : i32} loc(unknown), %arg249: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 249 : i32} loc(unknown), %arg250: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 250 : i32} loc(unknown), %arg251: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 251 : i32} loc(unknown), %arg252: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 252 : i32} loc(unknown), %arg253: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 253 : i32} loc(unknown), %arg254: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 254 : i32} loc(unknown), %arg255: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 255 : i32} loc(unknown), %arg256: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 256 : i32} loc(unknown), %arg257: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 257 : i32} loc(unknown), %arg258: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 258 : i32} loc(unknown), %arg259: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 259 : i32} loc(unknown), %arg260: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 260 : i32} loc(unknown), %arg261: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 261 : i32} loc(unknown), %arg262: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 262 : i32} loc(unknown), %arg263: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 263 : i32} loc(unknown), %arg264: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 264 : i32} loc(unknown), %arg265: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 265 : i32} loc(unknown), %arg266: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 266 : i32} loc(unknown), %arg267: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 267 : i32} loc(unknown), %arg268: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 268 : i32} loc(unknown), %arg269: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 269 : i32} loc(unknown), %arg270: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 270 : i32} loc(unknown), %arg271: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 271 : i32} loc(unknown), %arg272: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 272 : i32} loc(unknown), %arg273: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 273 : i32} loc(unknown), %arg274: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 274 : i32} loc(unknown), %arg275: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 275 : i32} loc(unknown), %arg276: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 276 : i32} loc(unknown), %arg277: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 277 : i32} loc(unknown), %arg278: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 278 : i32} loc(unknown), %arg279: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 279 : i32} loc(unknown), %arg280: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 280 : i32} loc(unknown), %arg281: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 281 : i32} loc(unknown), %arg282: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 282 : i32} loc(unknown), %arg283: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 283 : i32} loc(unknown), %arg284: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 284 : i32} loc(unknown), %arg285: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 285 : i32} loc(unknown), %arg286: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 286 : i32} loc(unknown), %arg287: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 287 : i32} loc(unknown), %arg288: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 288 : i32} loc(unknown), %arg289: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 289 : i32} loc(unknown), %arg290: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 290 : i32} loc(unknown), %arg291: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 291 : i32} loc(unknown), %arg292: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 292 : i32} loc(unknown), %arg293: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 293 : i32} loc(unknown), %arg294: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 294 : i32} loc(unknown), %arg295: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 295 : i32} loc(unknown), %arg296: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 296 : i32} loc(unknown), %arg297: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 297 : i32} loc(unknown), %arg298: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 298 : i32} loc(unknown), %arg299: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 299 : i32} loc(unknown), %arg300: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 300 : i32} loc(unknown), %arg301: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 301 : i32} loc(unknown), %arg302: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 302 : i32} loc(unknown), %arg303: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 303 : i32} loc(unknown), %arg304: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 304 : i32} loc(unknown), %arg305: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 305 : i32} loc(unknown), %arg306: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 306 : i32} loc(unknown), %arg307: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 307 : i32} loc(unknown), %arg308: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 308 : i32} loc(unknown), %arg309: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 309 : i32} loc(unknown), %arg310: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 310 : i32} loc(unknown), %arg311: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 311 : i32} loc(unknown), %arg312: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 312 : i32} loc(unknown), %arg313: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 313 : i32} loc(unknown), %arg314: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 314 : i32} loc(unknown), %arg315: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 315 : i32} loc(unknown), %arg316: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 316 : i32} loc(unknown), %arg317: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 317 : i32} loc(unknown), %arg318: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 318 : i32} loc(unknown), %arg319: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 319 : i32} loc(unknown), %arg320: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 320 : i32} loc(unknown), %arg321: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 321 : i32} loc(unknown), %arg322: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 322 : i32} loc(unknown), %arg323: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 323 : i32} loc(unknown), %arg324: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 324 : i32} loc(unknown), %arg325: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 325 : i32} loc(unknown), %arg326: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 326 : i32} loc(unknown), %arg327: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 327 : i32} loc(unknown), %arg328: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 328 : i32} loc(unknown), %arg329: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 329 : i32} loc(unknown), %arg330: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 330 : i32} loc(unknown), %arg331: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 331 : i32} loc(unknown), %arg332: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 332 : i32} loc(unknown), %arg333: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 333 : i32} loc(unknown), %arg334: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 334 : i32} loc(unknown), %arg335: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 335 : i32} loc(unknown), %arg336: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 336 : i32} loc(unknown), %arg337: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 337 : i32} loc(unknown), %arg338: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 338 : i32} loc(unknown), %arg339: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 339 : i32} loc(unknown), %arg340: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 340 : i32} loc(unknown), %arg341: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 341 : i32} loc(unknown), %arg342: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 342 : i32} loc(unknown), %arg343: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 343 : i32} loc(unknown), %arg344: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 344 : i32} loc(unknown), %arg345: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 345 : i32} loc(unknown), %arg346: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 346 : i32} loc(unknown), %arg347: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 347 : i32} loc(unknown), %arg348: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 348 : i32} loc(unknown), %arg349: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 349 : i32} loc(unknown), %arg350: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 350 : i32} loc(unknown), %arg351: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 351 : i32} loc(unknown), %arg352: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 352 : i32} loc(unknown), %arg353: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 353 : i32} loc(unknown), %arg354: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 354 : i32} loc(unknown), %arg355: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 355 : i32} loc(unknown), %arg356: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 356 : i32} loc(unknown), %arg357: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 357 : i32} loc(unknown), %arg358: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 358 : i32} loc(unknown), %arg359: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 359 : i32} loc(unknown), %arg360: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 360 : i32} loc(unknown), %arg361: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 361 : i32} loc(unknown), %arg362: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 362 : i32} loc(unknown), %arg363: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 363 : i32} loc(unknown), %arg364: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 364 : i32} loc(unknown), %arg365: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 365 : i32} loc(unknown), %arg366: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 366 : i32} loc(unknown), %arg367: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 367 : i32} loc(unknown), %arg368: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 368 : i32} loc(unknown), %arg369: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 369 : i32} loc(unknown), %arg370: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 370 : i32} loc(unknown), %arg371: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 371 : i32} loc(unknown), %arg372: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 372 : i32} loc(unknown), %arg373: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 373 : i32} loc(unknown), %arg374: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 374 : i32} loc(unknown), %arg375: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 375 : i32} loc(unknown), %arg376: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 376 : i32} loc(unknown), %arg377: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 377 : i32} loc(unknown), %arg378: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 378 : i32} loc(unknown), %arg379: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 379 : i32} loc(unknown), %arg380: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 380 : i32} loc(unknown), %arg381: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 381 : i32} loc(unknown), %arg382: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 382 : i32} loc(unknown), %arg383: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 383 : i32} loc(unknown), %arg384: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 384 : i32} loc(unknown), %arg385: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 385 : i32} loc(unknown), %arg386: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 386 : i32} loc(unknown), %arg387: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 387 : i32} loc(unknown), %arg388: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 388 : i32} loc(unknown), %arg389: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 389 : i32} loc(unknown), %arg390: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 390 : i32} loc(unknown), %arg391: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 391 : i32} loc(unknown), %arg392: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 392 : i32} loc(unknown), %arg393: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 393 : i32} loc(unknown), %arg394: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 394 : i32} loc(unknown), %arg395: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 395 : i32} loc(unknown), %arg396: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 396 : i32} loc(unknown), %arg397: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 397 : i32} loc(unknown), %arg398: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 398 : i32} loc(unknown), %arg399: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 399 : i32} loc(unknown), %arg400: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 400 : i32} loc(unknown), %arg401: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 401 : i32} loc(unknown), %arg402: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 402 : i32} loc(unknown), %arg403: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 403 : i32} loc(unknown), %arg404: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 404 : i32} loc(unknown), %arg405: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 405 : i32} loc(unknown), %arg406: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 406 : i32} loc(unknown), %arg407: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 407 : i32} loc(unknown), %arg408: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 408 : i32} loc(unknown), %arg409: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 409 : i32} loc(unknown), %arg410: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 410 : i32} loc(unknown), %arg411: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 411 : i32} loc(unknown), %arg412: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 412 : i32} loc(unknown), %arg413: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 413 : i32} loc(unknown), %arg414: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 414 : i32} loc(unknown), %arg415: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 415 : i32} loc(unknown), %arg416: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 416 : i32} loc(unknown), %arg417: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 417 : i32} loc(unknown), %arg418: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 418 : i32} loc(unknown), %arg419: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 419 : i32} loc(unknown), %arg420: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 420 : i32} loc(unknown), %arg421: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 421 : i32} loc(unknown), %arg422: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 422 : i32} loc(unknown), %arg423: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 423 : i32} loc(unknown), %arg424: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 424 : i32} loc(unknown), %arg425: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 425 : i32} loc(unknown), %arg426: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 426 : i32} loc(unknown), %arg427: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 427 : i32} loc(unknown), %arg428: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 428 : i32} loc(unknown), %arg429: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 429 : i32} loc(unknown), %arg430: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 430 : i32} loc(unknown), %arg431: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 431 : i32} loc(unknown), %arg432: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 432 : i32} loc(unknown), %arg433: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 433 : i32} loc(unknown), %arg434: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 434 : i32} loc(unknown), %arg435: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 435 : i32} loc(unknown), %arg436: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 436 : i32} loc(unknown), %arg437: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 437 : i32} loc(unknown), %arg438: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 438 : i32} loc(unknown), %arg439: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 439 : i32} loc(unknown), %arg440: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 440 : i32} loc(unknown), %arg441: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 441 : i32} loc(unknown), %arg442: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 442 : i32} loc(unknown), %arg443: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 443 : i32} loc(unknown), %arg444: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 444 : i32} loc(unknown), %arg445: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 445 : i32} loc(unknown), %arg446: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 446 : i32} loc(unknown), %arg447: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 447 : i32} loc(unknown), %arg448: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 448 : i32} loc(unknown), %arg449: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 449 : i32} loc(unknown), %arg450: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 450 : i32} loc(unknown), %arg451: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 451 : i32} loc(unknown), %arg452: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 452 : i32} loc(unknown), %arg453: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 453 : i32} loc(unknown), %arg454: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 454 : i32} loc(unknown), %arg455: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 455 : i32} loc(unknown), %arg456: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 456 : i32} loc(unknown), %arg457: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 457 : i32} loc(unknown), %arg458: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 458 : i32} loc(unknown), %arg459: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 459 : i32} loc(unknown), %arg460: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 460 : i32} loc(unknown), %arg461: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 461 : i32} loc(unknown), %arg462: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 462 : i32} loc(unknown), %arg463: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 463 : i32} loc(unknown), %arg464: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 464 : i32} loc(unknown), %arg465: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 465 : i32} loc(unknown), %arg466: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 466 : i32} loc(unknown), %arg467: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 467 : i32} loc(unknown), %arg468: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 468 : i32} loc(unknown), %arg469: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 469 : i32} loc(unknown), %arg470: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 470 : i32} loc(unknown), %arg471: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 471 : i32} loc(unknown), %arg472: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 472 : i32} loc(unknown), %arg473: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 473 : i32} loc(unknown), %arg474: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 474 : i32} loc(unknown), %arg475: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 475 : i32} loc(unknown), %arg476: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 476 : i32} loc(unknown), %arg477: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 477 : i32} loc(unknown), %arg478: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 478 : i32} loc(unknown), %arg479: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 479 : i32} loc(unknown), %arg480: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 480 : i32} loc(unknown), %arg481: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 481 : i32} loc(unknown), %arg482: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 482 : i32} loc(unknown), %arg483: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 483 : i32} loc(unknown), %arg484: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 484 : i32} loc(unknown), %arg485: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 485 : i32} loc(unknown), %arg486: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 486 : i32} loc(unknown), %arg487: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 487 : i32} loc(unknown), %arg488: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 488 : i32} loc(unknown), %arg489: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 489 : i32} loc(unknown), %arg490: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 490 : i32} loc(unknown), %arg491: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 491 : i32} loc(unknown), %arg492: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 492 : i32} loc(unknown), %arg493: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 493 : i32} loc(unknown), %arg494: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 494 : i32} loc(unknown), %arg495: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 495 : i32} loc(unknown), %arg496: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 496 : i32} loc(unknown), %arg497: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 497 : i32} loc(unknown), %arg498: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 498 : i32} loc(unknown), %arg499: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 499 : i32} loc(unknown), %arg500: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 500 : i32} loc(unknown), %arg501: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 501 : i32} loc(unknown), %arg502: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 502 : i32} loc(unknown), %arg503: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 503 : i32} loc(unknown), %arg504: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 504 : i32} loc(unknown), %arg505: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 505 : i32} loc(unknown), %arg506: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 506 : i32} loc(unknown), %arg507: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 507 : i32} loc(unknown), %arg508: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 508 : i32} loc(unknown), %arg509: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 509 : i32} loc(unknown), %arg510: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 510 : i32} loc(unknown), %arg511: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 511 : i32} loc(unknown), %arg512: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 512 : i32} loc(unknown), %arg513: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 513 : i32} loc(unknown), %arg514: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 514 : i32} loc(unknown), %arg515: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 515 : i32} loc(unknown), %arg516: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 516 : i32} loc(unknown), %arg517: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 517 : i32} loc(unknown), %arg518: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 518 : i32} loc(unknown), %arg519: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 519 : i32} loc(unknown), %arg520: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 520 : i32} loc(unknown), %arg521: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 521 : i32} loc(unknown), %arg522: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 522 : i32} loc(unknown), %arg523: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 523 : i32} loc(unknown), %arg524: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 524 : i32} loc(unknown), %arg525: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 525 : i32} loc(unknown), %arg526: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 526 : i32} loc(unknown), %arg527: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 527 : i32} loc(unknown), %arg528: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 528 : i32} loc(unknown), %arg529: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 529 : i32} loc(unknown), %arg530: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 530 : i32} loc(unknown), %arg531: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 531 : i32} loc(unknown), %arg532: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 532 : i32} loc(unknown), %arg533: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 533 : i32} loc(unknown), %arg534: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 534 : i32} loc(unknown), %arg535: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 535 : i32} loc(unknown), %arg536: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 536 : i32} loc(unknown), %arg537: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 537 : i32} loc(unknown), %arg538: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 538 : i32} loc(unknown), %arg539: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 539 : i32} loc(unknown), %arg540: tensor<1x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 540 : i32} loc(unknown), %arg541: tensor<1x2048x8192xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 541 : i32} loc(unknown), %arg542: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 542 : i32} loc(unknown), %arg543: tensor<1x8192x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 543 : i32} loc(unknown), %arg544: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 544 : i32} loc(unknown), %arg545: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 545 : i32} loc(unknown), %arg546: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 546 : i32} loc(unknown), %arg547: tensor<1x2048xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 547 : i32} loc(unknown), %arg548: tensor<1x3x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 548 : i32} loc(unknown), %arg549: tensor<1x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 549 : i32} loc(unknown), %arg550: tensor<1x2048x32x64xf32> {jax.arg_info = "state.opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 550 : i32} loc(unknown), %arg551: tensor<1xi32> {jax.arg_info = "state.opt_states[0][<flat index 1>][3][<flat index 0>]", mhlo.sharding = "{maximal device=0}", tf.aliasing_output = 551 : i32} loc(unknown), %arg552: tensor<1x2048xi32> {jax.arg_info = "inputs[<flat index 1>]", mhlo.sharding = "{maximal device=0}"} loc(unknown), %arg553: tensor<1x2048xi32> {jax.arg_info = "inputs[<flat index 2>]", mhlo.sharding = "{maximal device=0}"} loc(unknown), %arg554: tensor<1x2048xi32> {jax.arg_info = "inputs[<flat index 3>]", mhlo.sharding = "{maximal device=0}"} loc(unknown), %arg555: tensor<1x2048xi32> {jax.arg_info = "inputs[<flat index 4>]", mhlo.sharding = "{maximal device=0}"} loc(unknown), %arg556: tensor<1x2048xi32> {jax.arg_info = "inputs[<flat index 5>]", mhlo.sharding = "{maximal device=0}"} loc(unknown), %arg557: tensor<1x2048xi32> {jax.arg_info = "inputs[<flat index 6>]", mhlo.sharding = "{maximal device=0}"} loc(unknown)) -> (tensor<ui32> {jax.result_info = "[0].step", mhlo.sharding = "{replicated}"}, tensor<2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['final_ln']['bias']", mhlo.sharding = "{replicated}"}, tensor<2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['final_ln']['scale']", mhlo.sharding = "{replicated}"}, tensor<2048x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['position_emb']['emb_var']", mhlo.sharding = "{maximal device=0}"}, tensor<51200xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['softmax']['logits_ffn']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<2048x51200xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['softmax']['logits_ffn']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].mdl_vars['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<i32> {jax.result_info = "[0].opt_states[0][<flat index 0>][0][<flat index 0>]", mhlo.sharding = "{replicated}"}, tensor<i32> {jax.result_info = "[0].opt_states[0][<flat index 0>][1][<flat index 0>]", mhlo.sharding = "{replicated}"}, tensor<i32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 0>]", mhlo.sharding = "{replicated}"}, tensor<2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['final_ln']['bias']", mhlo.sharding = "{replicated}"}, tensor<2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['final_ln']['scale']", mhlo.sharding = "{replicated}"}, tensor<2048x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['position_emb']['emb_var']", mhlo.sharding = "{maximal device=0}"}, tensor<51200xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['softmax']['logits_ffn']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<2048x51200xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 1>]['params']['lm']['softmax']['logits_ffn']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['final_ln']['bias']", mhlo.sharding = "{replicated}"}, tensor<2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['final_ln']['scale']", mhlo.sharding = "{replicated}"}, tensor<2048x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['position_emb']['emb_var']", mhlo.sharding = "{maximal device=0}"}, tensor<51200xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['softmax']['logits_ffn']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<2048x51200xf32> {jax.result_info = "[0].opt_states[0][<flat index 0>][2][<flat index 2>]['params']['lm']['softmax']['logits_ffn']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<i32> {jax.result_info = "[0].opt_states[0][<flat index 0>][3][<flat index 0>]", mhlo.sharding = "{replicated}"}, tensor<1xi32> {jax.result_info = "[0].opt_states[0][<flat index 1>][0][<flat index 0>]", mhlo.sharding = "{maximal device=0}"}, tensor<1xi32> {jax.result_info = "[0].opt_states[0][<flat index 1>][1][<flat index 0>]", mhlo.sharding = "{maximal device=0}"}, tensor<1xi32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 0>]", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 1>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_0']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_1']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_10']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_11']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_12']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_13']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_14']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_15']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_2']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_3']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_4']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_5']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_6']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_7']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_8']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x8192xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer1']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['bias']['b']", mhlo.sharding = "{maximal device=0}"}, tensor<1x8192x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['ffn_layer2']['linear']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['ff_layer']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['bias']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['layer_norm']['scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x3x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['combined_qkv']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['per_dim_scale']['per_dim_scale']", mhlo.sharding = "{maximal device=0}"}, tensor<1x2048x32x64xf32> {jax.result_info = "[0].opt_states[0][<flat index 1>][2][<flat index 2>]['params']['lm']['transformer']['pipeline']['body']['x_layers_9']['self_attention']['post']['w']", mhlo.sharding = "{maximal device=0}"}, tensor<1xi32> {jax.result_info = "[0].opt_states[0][<flat index 1>][3][<flat index 0>]", mhlo.sharding = "{maximal device=0}"}, tensor<f32> {jax.result_info = "[1]", mhlo.sharding = "{replicated}"}, tensor<bf16> {jax.result_info = "[2][<flat index 0>][0]", mhlo.sharding = "{replicated}"}, tensor<bf16> {jax.result_info = "[2][<flat index 0>][1]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 1>][0]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 1>][1]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 2>][0]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 2>][1]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 3>][0]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 3>][1]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 4>][0]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 4>][1]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 5>][0]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[2][<flat index 5>][1]", mhlo.sharding = "{replicated}"}, tensor<1x2048xi32> {jax.result_info = "[3][<flat index 0>]", mhlo.sharding = "{replicated}"}, tensor<1xf32> {jax.result_info = "[3][<flat index 1>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 0>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 1>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 2>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 3>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 4>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 5>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 6>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 7>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 8>]", mhlo.sharding = "{replicated}"}, tensor<bf16> {jax.result_info = "[4][<flat index 9>]", mhlo.sharding = "{replicated}"}, tensor<bf16> {jax.result_info = "[4][<flat index 10>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 11>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 12>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 13>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 14>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 15>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 16>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 17>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 18>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 19>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 20>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 21>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 22>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 23>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 24>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 25>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 26>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 27>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 28>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 29>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 30>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 31>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 32>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 33>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 34>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 35>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 36>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 37>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 38>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 39>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 40>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 41>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 42>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 43>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 44>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 45>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 46>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 47>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 48>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 49>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 50>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 51>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 52>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 53>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 54>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 55>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 56>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 57>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 58>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 59>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 60>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 61>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 62>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 63>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 64>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 65>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 66>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 67>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 68>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 69>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 70>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 71>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 72>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 73>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 74>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 75>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 76>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 77>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 78>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 79>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 80>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 81>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 82>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 83>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 84>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 85>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 86>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 87>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 88>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 89>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 90>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 91>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 92>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 93>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 94>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 95>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 96>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 97>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 98>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 99>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 100>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 101>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 102>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 103>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 104>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 105>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 106>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 107>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 108>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 109>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 110>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 111>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 112>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 113>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 114>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 115>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 116>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 117>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 118>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 119>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 120>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 121>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 122>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 123>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 124>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 125>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 126>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 127>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 128>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 129>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 130>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 131>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 132>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 133>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 134>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 135>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 136>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 137>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 138>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 139>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 140>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 141>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 142>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 143>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 144>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 145>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 146>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 147>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 148>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 149>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 150>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 151>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 152>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 153>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 154>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 155>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 156>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 157>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 158>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 159>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 160>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 161>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 162>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 163>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 164>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 165>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 166>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 167>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 168>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 169>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 170>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 171>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 172>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 173>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 174>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 175>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 176>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 177>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 178>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 179>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 180>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 181>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 182>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 183>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 184>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 185>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 186>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 187>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 188>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 189>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 190>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 191>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 192>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 193>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 194>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 195>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 196>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 197>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 198>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 199>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 200>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 201>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 202>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 203>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 204>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 205>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 206>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 207>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 208>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 209>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 210>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 211>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 212>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 213>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 214>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 215>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 216>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 217>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 218>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 219>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 220>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 221>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 222>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 223>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 224>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 225>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 226>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 227>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 228>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 229>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 230>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 231>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 232>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 233>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 234>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 235>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 236>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 237>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 238>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 239>]", mhlo.sharding = "{replicated}"}, tensor<f32> {jax.result_info = "[4][<flat index 240>]", mhlo.sharding = "{replicated}"}) {
    %ptrue = arith.constant 1 : i1
    %0 = arith.select %ptrue, %arg552, %arg552 : tensor<1x2048xi32>
    %1 = arith.select %ptrue, %arg553, %arg553 : tensor<1x2048xi32>
    %2 = arith.select %ptrue, %arg554, %arg554 : tensor<1x2048xi32>
    %3 = arith.select %ptrue, %arg555, %arg555 : tensor<1x2048xi32>
    %4 = arith.select %ptrue, %arg556, %arg556 : tensor<1x2048xi32>
    %5 = arith.select %ptrue, %arg557, %arg557 : tensor<1x2048xi32>
    %6 = arith.select %ptrue, %0, %0 : tensor<1x2048xi32>
    %7 = arith.select %ptrue, %1, %1 : tensor<1x2048xi32>
    %8 = arith.select %ptrue, %2, %2 : tensor<1x2048xi32>
    %9 = arith.select %ptrue, %3, %3 : tensor<1x2048xi32>
    %10 = arith.select %ptrue, %4, %4 : tensor<1x2048xi32>
    %11 = arith.select %ptrue, %5, %5 : tensor<1x2048xi32>
    %12 = stablehlo.convert %arg1 : (tensor<2048xf32>) -> tensor<2048xbf16> loc(#loc135)
    %13 = stablehlo.convert %arg2 : (tensor<2048xf32>) -> tensor<2048xbf16> loc(#loc135)
    %14 = stablehlo.convert %arg3 : (tensor<2048x2048xf32>) -> tensor<2048x2048xbf16> loc(#loc135)
    %15 = stablehlo.convert %arg4 : (tensor<51200xf32>) -> tensor<51200xbf16> loc(#loc135)
    %16 = stablehlo.convert %arg5 : (tensor<2048x51200xf32>) -> tensor<2048x51200xbf16> loc(#loc135)
    %17 = stablehlo.convert %arg6 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %18 = stablehlo.convert %arg7 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %19 = stablehlo.convert %arg8 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %20 = stablehlo.convert %arg9 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %21 = stablehlo.convert %arg10 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %22 = stablehlo.convert %arg11 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %23 = stablehlo.convert %arg12 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %24 = stablehlo.convert %arg13 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %25 = stablehlo.convert %arg14 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %26 = stablehlo.convert %arg15 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %27 = stablehlo.convert %arg16 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %28 = stablehlo.convert %arg17 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %29 = stablehlo.convert %arg18 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %30 = stablehlo.convert %arg19 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %31 = stablehlo.convert %arg20 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %32 = stablehlo.convert %arg21 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %33 = stablehlo.convert %arg22 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %34 = stablehlo.convert %arg23 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %35 = stablehlo.convert %arg24 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %36 = stablehlo.convert %arg25 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %37 = stablehlo.convert %arg26 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %38 = stablehlo.convert %arg27 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %39 = stablehlo.convert %arg28 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %40 = stablehlo.convert %arg29 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %41 = stablehlo.convert %arg30 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %42 = stablehlo.convert %arg31 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %43 = stablehlo.convert %arg32 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %44 = stablehlo.convert %arg33 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %45 = stablehlo.convert %arg34 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %46 = stablehlo.convert %arg35 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %47 = stablehlo.convert %arg36 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %48 = stablehlo.convert %arg37 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %49 = stablehlo.convert %arg38 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %50 = stablehlo.convert %arg39 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %51 = stablehlo.convert %arg40 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %52 = stablehlo.convert %arg41 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %53 = stablehlo.convert %arg42 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %54 = stablehlo.convert %arg43 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %55 = stablehlo.convert %arg44 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %56 = stablehlo.convert %arg45 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %57 = stablehlo.convert %arg46 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %58 = stablehlo.convert %arg47 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %59 = stablehlo.convert %arg48 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %60 = stablehlo.convert %arg49 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %61 = stablehlo.convert %arg50 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %62 = stablehlo.convert %arg51 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %63 = stablehlo.convert %arg52 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %64 = stablehlo.convert %arg53 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %65 = stablehlo.convert %arg54 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %66 = stablehlo.convert %arg55 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %67 = stablehlo.convert %arg56 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %68 = stablehlo.convert %arg57 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %69 = stablehlo.convert %arg58 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %70 = stablehlo.convert %arg59 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %71 = stablehlo.convert %arg60 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %72 = stablehlo.convert %arg61 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %73 = stablehlo.convert %arg62 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %74 = stablehlo.convert %arg63 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %75 = stablehlo.convert %arg64 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %76 = stablehlo.convert %arg65 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %77 = stablehlo.convert %arg66 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %78 = stablehlo.convert %arg67 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %79 = stablehlo.convert %arg68 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %80 = stablehlo.convert %arg69 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %81 = stablehlo.convert %arg70 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %82 = stablehlo.convert %arg71 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %83 = stablehlo.convert %arg72 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %84 = stablehlo.convert %arg73 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %85 = stablehlo.convert %arg74 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %86 = stablehlo.convert %arg75 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %87 = stablehlo.convert %arg76 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %88 = stablehlo.convert %arg77 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %89 = stablehlo.convert %arg78 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %90 = stablehlo.convert %arg79 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %91 = stablehlo.convert %arg80 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %92 = stablehlo.convert %arg81 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %93 = stablehlo.convert %arg82 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %94 = stablehlo.convert %arg83 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %95 = stablehlo.convert %arg84 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %96 = stablehlo.convert %arg85 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %97 = stablehlo.convert %arg86 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %98 = stablehlo.convert %arg87 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %99 = stablehlo.convert %arg88 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %100 = stablehlo.convert %arg89 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %101 = stablehlo.convert %arg90 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %102 = stablehlo.convert %arg91 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %103 = stablehlo.convert %arg92 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %104 = stablehlo.convert %arg93 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %105 = stablehlo.convert %arg94 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %106 = stablehlo.convert %arg95 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %107 = stablehlo.convert %arg96 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %108 = stablehlo.convert %arg97 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %109 = stablehlo.convert %arg98 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %110 = stablehlo.convert %arg99 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %111 = stablehlo.convert %arg100 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %112 = stablehlo.convert %arg101 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %113 = stablehlo.convert %arg102 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %114 = stablehlo.convert %arg103 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %115 = stablehlo.convert %arg104 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %116 = stablehlo.convert %arg105 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %117 = stablehlo.convert %arg106 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %118 = stablehlo.convert %arg107 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %119 = stablehlo.convert %arg108 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %120 = stablehlo.convert %arg109 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %121 = stablehlo.convert %arg110 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %122 = stablehlo.convert %arg111 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %123 = stablehlo.convert %arg112 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %124 = stablehlo.convert %arg113 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %125 = stablehlo.convert %arg114 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %126 = stablehlo.convert %arg115 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %127 = stablehlo.convert %arg116 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %128 = stablehlo.convert %arg117 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %129 = stablehlo.convert %arg118 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %130 = stablehlo.convert %arg119 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %131 = stablehlo.convert %arg120 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %132 = stablehlo.convert %arg121 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %133 = stablehlo.convert %arg122 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %134 = stablehlo.convert %arg123 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %135 = stablehlo.convert %arg124 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %136 = stablehlo.convert %arg125 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %137 = stablehlo.convert %arg126 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %138 = stablehlo.convert %arg127 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %139 = stablehlo.convert %arg128 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %140 = stablehlo.convert %arg129 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %141 = stablehlo.convert %arg130 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %142 = stablehlo.convert %arg131 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %143 = stablehlo.convert %arg132 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %144 = stablehlo.convert %arg133 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %145 = stablehlo.convert %arg134 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %146 = stablehlo.convert %arg135 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %147 = stablehlo.convert %arg136 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %148 = stablehlo.convert %arg137 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %149 = stablehlo.convert %arg138 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %150 = stablehlo.convert %arg139 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %151 = stablehlo.convert %arg140 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %152 = stablehlo.convert %arg141 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %153 = stablehlo.convert %arg142 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %154 = stablehlo.convert %arg143 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %155 = stablehlo.convert %arg144 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %156 = stablehlo.convert %arg145 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %157 = stablehlo.convert %arg146 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %158 = stablehlo.convert %arg147 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %159 = stablehlo.convert %arg148 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %160 = stablehlo.convert %arg149 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %161 = stablehlo.convert %arg150 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %162 = stablehlo.convert %arg151 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %163 = stablehlo.convert %arg152 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %164 = stablehlo.convert %arg153 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %165 = stablehlo.convert %arg154 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %166 = stablehlo.convert %arg155 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %167 = stablehlo.convert %arg156 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %168 = stablehlo.convert %arg157 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %169 = stablehlo.convert %arg158 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %170 = stablehlo.convert %arg159 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %171 = stablehlo.convert %arg160 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %172 = stablehlo.convert %arg161 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %173 = stablehlo.convert %arg162 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %174 = stablehlo.convert %arg163 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %175 = stablehlo.convert %arg164 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %176 = stablehlo.convert %arg165 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %177 = stablehlo.convert %arg166 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %178 = stablehlo.convert %arg167 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %179 = stablehlo.convert %arg168 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %180 = stablehlo.convert %arg169 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %181 = stablehlo.convert %arg170 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %182 = stablehlo.convert %arg171 : (tensor<1x8192xf32>) -> tensor<1x8192xbf16> loc(#loc135)
    %183 = stablehlo.convert %arg172 : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xbf16> loc(#loc135)
    %184 = stablehlo.convert %arg173 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %185 = stablehlo.convert %arg174 : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xbf16> loc(#loc135)
    %186 = stablehlo.convert %arg175 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %187 = stablehlo.convert %arg176 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %188 = stablehlo.convert %arg177 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %189 = stablehlo.convert %arg178 : (tensor<1x2048xf32>) -> tensor<1x2048xbf16> loc(#loc135)
    %190 = stablehlo.convert %arg179 : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xbf16> loc(#loc135)
    %191 = stablehlo.convert %arg180 : (tensor<1x64xf32>) -> tensor<1x64xbf16> loc(#loc135)
    %192 = stablehlo.convert %arg181 : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xbf16> loc(#loc135)
    %193 = stablehlo.convert %8 : (tensor<1x2048xi32>) -> tensor<1x2048xf32> loc(#loc136)
    %194 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %195 = stablehlo.broadcast_in_dim %194, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc137)
    %196 = stablehlo.subtract %195, %193 : tensor<1x2048xf32> loc(#loc137)
    %197 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc138)
    %198 = stablehlo.reduce(%196 init: %197) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc138)
      stablehlo.return %29929 : tensor<f32> loc(#loc138)
    } loc(#loc138)
    %199 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %200 = stablehlo.divide %198, %199 : tensor<f32> loc(#loc139)
    %201 = stablehlo.transpose %16, dims = [1, 0] : (tensor<2048x51200xbf16>) -> tensor<51200x2048xbf16> loc(#loc140)
    %202 = call @_one_hot(%6) : (tensor<1x2048xi32>) -> tensor<1x2048x51200xbf16> loc(#loc141)
    %203 = call @_einsum(%202, %201) : (tensor<1x2048x51200xbf16>, tensor<51200x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc142)
    %204 = stablehlo.constant dense<4.525000e+01> : tensor<bf16> loc(#loc)
    %205 = stablehlo.broadcast_in_dim %204, dims = [] : (tensor<bf16>) -> tensor<1x2048x2048xbf16> loc(#loc143)
    %206 = stablehlo.multiply %203, %205 : tensor<1x2048x2048xbf16> loc(#loc143)
    %207 = arith.select %ptrue, %206, %206 : tensor<1x2048x2048xbf16>
    %208 = "stablehlo.slice"(%14) {limit_indices = dense<2048> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2048x2048xbf16>) -> tensor<2048x2048xbf16> loc(#loc145)
    %209 = call @_one_hot_0(%10) : (tensor<1x2048xi32>) -> tensor<1x2048x2048xbf16> loc(#loc146)
    %210 = call @_einsum_1(%209, %208) : (tensor<1x2048x2048xbf16>, tensor<2048x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc147)
    %211 = arith.select %ptrue, %210, %210 : tensor<1x2048x2048xbf16>
    %212 = stablehlo.add %207, %211 : tensor<1x2048x2048xbf16> loc(#loc149)
    %213 = stablehlo.broadcast_in_dim %9, dims = [0, 1] : (tensor<1x2048xi32>) -> tensor<1x2048x1xi32> loc(#loc150)
    %214 = stablehlo.broadcast_in_dim %9, dims = [0, 2] : (tensor<1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc151)
    %215 = stablehlo.broadcast_in_dim %213, dims = [0, 1, 2] : (tensor<1x2048x1xi32>) -> tensor<1x2048x2048xi32> loc(#loc152)
    %216 = stablehlo.broadcast_in_dim %214, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x2048x2048xi32> loc(#loc152)
    %217 = stablehlo.compare  NE, %215, %216,  SIGNED : (tensor<1x2048x2048xi32>, tensor<1x2048x2048xi32>) -> tensor<1x2048x2048xi1> loc(#loc152)
    %218 = stablehlo.convert %217 : (tensor<1x2048x2048xi1>) -> tensor<1x2048x2048xbf16> loc(#loc153)
    %219 = stablehlo.broadcast_in_dim %218, dims = [0, 2, 3] : (tensor<1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc154)
    %220 = stablehlo.constant dense<-2.366030e+38> : tensor<bf16> loc(#loc)
    %221 = stablehlo.broadcast_in_dim %220, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc155)
    %222 = stablehlo.multiply %219, %221 : tensor<1x1x2048x2048xbf16> loc(#loc155)
    %223 = stablehlo.iota dim = 0 : tensor<2048xi32> loc(#loc156)
    %224 = stablehlo.broadcast_in_dim %223, dims = [1] : (tensor<2048xi32>) -> tensor<1x2048xi32> loc(#loc157)
    %225 = stablehlo.reshape %224 : (tensor<1x2048xi32>) -> tensor<1x1x1x2048xi32> loc(#loc158)
    %226 = stablehlo.reshape %225 : (tensor<1x1x1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc159)
    %227 = stablehlo.broadcast_in_dim %226, dims = [1, 2, 3] : (tensor<1x1x2048xi32>) -> tensor<2048x1x1x2048xi32> loc(#loc160)
    %228 = stablehlo.reshape %227 : (tensor<2048x1x1x2048xi32>) -> tensor<2048x2048xi32> loc(#loc161)
    %229 = stablehlo.iota dim = 0 : tensor<2048xi32> loc(#loc162)
    %230 = stablehlo.broadcast_in_dim %229, dims = [0] : (tensor<2048xi32>) -> tensor<2048x1xi32> loc(#loc163)
    %231 = stablehlo.reshape %230 : (tensor<2048x1xi32>) -> tensor<1x2048x1x1xi32> loc(#loc164)
    %232 = stablehlo.reshape %231 : (tensor<1x2048x1x1xi32>) -> tensor<1x2048x1xi32> loc(#loc165)
    %233 = stablehlo.broadcast_in_dim %232, dims = [0, 1, 3] : (tensor<1x2048x1xi32>) -> tensor<1x2048x2048x1xi32> loc(#loc166)
    %234 = stablehlo.reshape %233 : (tensor<1x2048x2048x1xi32>) -> tensor<2048x2048xi32> loc(#loc167)
    %235 = stablehlo.compare  LT, %234, %228,  SIGNED : (tensor<2048x2048xi32>, tensor<2048x2048xi32>) -> tensor<2048x2048xi1> loc(#loc168)
    %236 = stablehlo.convert %235 : (tensor<2048x2048xi1>) -> tensor<2048x2048xbf16> loc(#loc169)
    %237 = stablehlo.constant dense<-2.366030e+38> : tensor<bf16> loc(#loc)
    %238 = stablehlo.broadcast_in_dim %237, dims = [] : (tensor<bf16>) -> tensor<2048x2048xbf16> loc(#loc170)
    %239 = stablehlo.multiply %236, %238 : tensor<2048x2048xbf16> loc(#loc170)
    %240 = stablehlo.broadcast_in_dim %239, dims = [2, 3] : (tensor<2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc171)
    %241 = stablehlo.minimum %222, %240 : tensor<1x1x2048x2048xbf16> loc(#loc172)
    %242 = arith.select %ptrue, %212, %212 : tensor<1x2048x2048xbf16>
    %243 = arith.select %ptrue, %8, %8 : tensor<1x2048xi32>
    %244 = arith.select %ptrue, %241, %241 : tensor<1x1x2048x2048xbf16>
    %245 = stablehlo.reshape %242 : (tensor<1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc175)
    %246 = stablehlo.transpose %245, dims = [0, 2, 1, 3, 4] : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc176)
    %247 = stablehlo.reshape %246 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc177)
    %248 = stablehlo.reshape %243 : (tensor<1x2048xi32>) -> tensor<1x1x1x2048xi32> loc(#loc178)
    %249 = stablehlo.transpose %248, dims = [0, 2, 1, 3] : (tensor<1x1x1x2048xi32>) -> tensor<1x1x1x2048xi32> loc(#loc179)
    %250 = stablehlo.reshape %249 : (tensor<1x1x1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc180)
    %251 = stablehlo.reshape %244 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x1x2048x2048xbf16> loc(#loc181)
    %252 = stablehlo.transpose %251, dims = [0, 2, 1, 3, 4, 5] : (tensor<1x1x1x1x2048x2048xbf16>) -> tensor<1x1x1x1x2048x2048xbf16> loc(#loc182)
    %253 = stablehlo.reshape %252 : (tensor<1x1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc183)
    %254 = stablehlo.broadcast_in_dim %247, dims = [0, 2, 3, 4] : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc184)
    %255 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %256 = call @_pad(%254, %255) : (tensor<1x1x1x2048x2048xbf16>, tensor<i32>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc185)
    %257 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %258 = stablehlo.broadcast_in_dim %257, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc186)
    %259 = stablehlo.iota dim = 0 : tensor<1xi32> loc(#loc187)
    %260 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %261:2 = call @remainder(%260) : (tensor<i32>) -> (tensor<i32>, tensor<i1>) loc(#loc188)
    %262 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %263:2 = call @remainder_2(%262) : (tensor<i32>) -> (tensor<i32>, tensor<i1>) loc(#loc189)
    %264 = stablehlo.iota dim = 0 : tensor<1x1x2048x2048xi32> loc(#loc190)
    %265 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %266 = stablehlo.broadcast_in_dim %265, dims = [] : (tensor<i32>) -> tensor<1x1x2048x2048xi32> loc(#loc191)
    %267 = stablehlo.compare  EQ, %264, %266,  SIGNED : (tensor<1x1x2048x2048xi32>, tensor<1x1x2048x2048xi32>) -> tensor<1x1x2048x2048xi1> loc(#loc191)
    %268 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %269 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %270 = stablehlo.add %268, %269 : tensor<i32> loc(#loc192)
    %271 = stablehlo.constant dense<false> : tensor<i1> loc(#loc)
    %272 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %273 = stablehlo.select %271, %270, %272 : tensor<i1>, tensor<i32> loc(#loc193)
    %274 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %275 = stablehlo.constant dense<2048> : tensor<i32> loc(#loc)
    %276 = stablehlo.add %274, %275 : tensor<i32> loc(#loc192)
    %277 = stablehlo.constant dense<false> : tensor<i1> loc(#loc)
    %278 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %279 = stablehlo.select %277, %276, %278 : tensor<i1>, tensor<i32> loc(#loc193)
    %280 = stablehlo.broadcast_in_dim %273, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc194)
    %281 = stablehlo.broadcast_in_dim %279, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc194)
    %282 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %283 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %284 = stablehlo.add %282, %283 : tensor<i32> loc(#loc192)
    %285 = stablehlo.constant dense<false> : tensor<i1> loc(#loc)
    %286 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %287 = stablehlo.select %285, %284, %286 : tensor<i1>, tensor<i32> loc(#loc193)
    %288 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %289 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %290 = stablehlo.add %288, %289 : tensor<i32> loc(#loc192)
    %291 = stablehlo.constant dense<false> : tensor<i1> loc(#loc)
    %292 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %293 = stablehlo.select %291, %290, %292 : tensor<i1>, tensor<i32> loc(#loc193)
    %294 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %295 = stablehlo.constant dense<2048> : tensor<i32> loc(#loc)
    %296 = stablehlo.add %294, %295 : tensor<i32> loc(#loc192)
    %297 = stablehlo.constant dense<false> : tensor<i1> loc(#loc)
    %298 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %299 = stablehlo.select %297, %296, %298 : tensor<i1>, tensor<i32> loc(#loc193)
    %300 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %301 = stablehlo.constant dense<2048> : tensor<i32> loc(#loc)
    %302 = stablehlo.add %300, %301 : tensor<i32> loc(#loc192)
    %303 = stablehlo.constant dense<false> : tensor<i1> loc(#loc)
    %304 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %305 = stablehlo.select %303, %302, %304 : tensor<i1>, tensor<i32> loc(#loc193)
    %306 = stablehlo.broadcast_in_dim %287, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc194)
    %307 = stablehlo.broadcast_in_dim %293, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc194)
    %308 = stablehlo.broadcast_in_dim %299, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc194)
    %309 = stablehlo.broadcast_in_dim %305, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc194)
    %310 = stablehlo.iota dim = 0 : tensor<1xi32> loc(#loc195)
    %311 = stablehlo.iota dim = 0 : tensor<2048xi32> loc(#loc196)
    %312 = stablehlo.broadcast_in_dim %311, dims = [1] : (tensor<2048xi32>) -> tensor<1x2048xi32> loc(#loc197)
    %313 = stablehlo.reshape %312 : (tensor<1x2048xi32>) -> tensor<1x1x1x2048xi32> loc(#loc198)
    %314 = stablehlo.reshape %313 : (tensor<1x1x1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc199)
    %315 = stablehlo.broadcast_in_dim %314, dims = [1, 2, 3] : (tensor<1x1x2048xi32>) -> tensor<2048x1x1x2048xi32> loc(#loc200)
    %316 = stablehlo.reshape %315 : (tensor<2048x1x1x2048xi32>) -> tensor<2048x2048xi32> loc(#loc201)
    %317 = stablehlo.iota dim = 0 : tensor<2048xi32> loc(#loc202)
    %318 = stablehlo.broadcast_in_dim %317, dims = [0] : (tensor<2048xi32>) -> tensor<2048x1xi32> loc(#loc203)
    %319 = stablehlo.reshape %318 : (tensor<2048x1xi32>) -> tensor<1x2048x1x1xi32> loc(#loc204)
    %320 = stablehlo.reshape %319 : (tensor<1x2048x1x1xi32>) -> tensor<1x2048x1xi32> loc(#loc205)
    %321 = stablehlo.broadcast_in_dim %320, dims = [0, 1, 3] : (tensor<1x2048x1xi32>) -> tensor<1x2048x2048x1xi32> loc(#loc206)
    %322 = stablehlo.reshape %321 : (tensor<1x2048x2048x1xi32>) -> tensor<2048x2048xi32> loc(#loc207)
    %323 = stablehlo.compare  LT, %322, %316,  SIGNED : (tensor<2048x2048xi32>, tensor<2048x2048xi32>) -> tensor<2048x2048xi1> loc(#loc208)
    %324 = stablehlo.convert %323 : (tensor<2048x2048xi1>) -> tensor<2048x2048xbf16> loc(#loc209)
    %325 = stablehlo.constant dense<-2.366030e+38> : tensor<bf16> loc(#loc)
    %326 = stablehlo.broadcast_in_dim %325, dims = [] : (tensor<bf16>) -> tensor<2048x2048xbf16> loc(#loc210)
    %327 = stablehlo.multiply %324, %326 : tensor<2048x2048xbf16> loc(#loc210)
    %328 = stablehlo.broadcast_in_dim %327, dims = [2, 3] : (tensor<2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc211)
    %329 = stablehlo.broadcast_in_dim %328, dims = [1, 2, 3, 4] : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc212)
    %330 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %331 = stablehlo.broadcast_in_dim %330, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc213)
    %332 = stablehlo.broadcast_in_dim %331, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc214)
    %333 = call @softplus() : () -> tensor<bf16> loc(#loc215)
    %334 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %335 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %336 = stablehlo.multiply %334, %335 : tensor<f32> loc(#loc216)
    %337 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %338 = call @_where_4(%337) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc217)
    %339 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %340 = stablehlo.broadcast_in_dim %339, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc218)
    %341 = stablehlo.broadcast_in_dim %340, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc219)
    %342 = call @softplus_5() : () -> tensor<bf16> loc(#loc220)
    %343 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %344 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %345 = stablehlo.multiply %343, %344 : tensor<f32> loc(#loc221)
    %346 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %347 = call @_where_6(%346) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc222)
    %348 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %349 = stablehlo.broadcast_in_dim %348, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc223)
    %350 = stablehlo.broadcast_in_dim %349, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc224)
    %351 = call @softplus_7() : () -> tensor<bf16> loc(#loc225)
    %352 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %353 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %354 = stablehlo.multiply %352, %353 : tensor<f32> loc(#loc226)
    %355 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %356 = call @_where_8(%355) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc227)
    %357 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %358 = stablehlo.broadcast_in_dim %357, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc228)
    %359 = stablehlo.broadcast_in_dim %358, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc229)
    %360 = call @softplus_9() : () -> tensor<bf16> loc(#loc230)
    %361 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %362 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %363 = stablehlo.multiply %361, %362 : tensor<f32> loc(#loc231)
    %364 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %365 = call @_where_10(%364) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc232)
    %366 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %367 = stablehlo.broadcast_in_dim %366, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc233)
    %368 = stablehlo.broadcast_in_dim %367, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc234)
    %369 = call @softplus_11() : () -> tensor<bf16> loc(#loc235)
    %370 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %371 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %372 = stablehlo.multiply %370, %371 : tensor<f32> loc(#loc236)
    %373 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %374 = call @_where_12(%373) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc237)
    %375 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %376 = stablehlo.broadcast_in_dim %375, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc238)
    %377 = stablehlo.broadcast_in_dim %376, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc239)
    %378 = call @softplus_13() : () -> tensor<bf16> loc(#loc240)
    %379 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %380 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %381 = stablehlo.multiply %379, %380 : tensor<f32> loc(#loc241)
    %382 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %383 = call @_where_14(%382) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc242)
    %384 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %385 = stablehlo.broadcast_in_dim %384, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc243)
    %386 = stablehlo.broadcast_in_dim %385, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc244)
    %387 = call @softplus_15() : () -> tensor<bf16> loc(#loc245)
    %388 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %389 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %390 = stablehlo.multiply %388, %389 : tensor<f32> loc(#loc246)
    %391 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %392 = call @_where_16(%391) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc247)
    %393 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %394 = stablehlo.broadcast_in_dim %393, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc248)
    %395 = stablehlo.broadcast_in_dim %394, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc249)
    %396 = call @softplus_17() : () -> tensor<bf16> loc(#loc250)
    %397 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %398 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %399 = stablehlo.multiply %397, %398 : tensor<f32> loc(#loc251)
    %400 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %401 = call @_where_18(%400) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc252)
    %402 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %403 = stablehlo.broadcast_in_dim %402, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc253)
    %404 = stablehlo.broadcast_in_dim %403, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc254)
    %405 = call @softplus_19() : () -> tensor<bf16> loc(#loc255)
    %406 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %407 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %408 = stablehlo.multiply %406, %407 : tensor<f32> loc(#loc256)
    %409 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %410 = call @_where_20(%409) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc257)
    %411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %412 = stablehlo.broadcast_in_dim %411, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc258)
    %413 = stablehlo.broadcast_in_dim %412, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc259)
    %414 = call @softplus_21() : () -> tensor<bf16> loc(#loc260)
    %415 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %416 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %417 = stablehlo.multiply %415, %416 : tensor<f32> loc(#loc261)
    %418 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %419 = call @_where_22(%418) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc262)
    %420 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %421 = stablehlo.broadcast_in_dim %420, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc263)
    %422 = stablehlo.broadcast_in_dim %421, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc264)
    %423 = call @softplus_23() : () -> tensor<bf16> loc(#loc265)
    %424 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %425 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %426 = stablehlo.multiply %424, %425 : tensor<f32> loc(#loc266)
    %427 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %428 = call @_where_24(%427) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc267)
    %429 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %430 = stablehlo.broadcast_in_dim %429, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc268)
    %431 = stablehlo.broadcast_in_dim %430, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc269)
    %432 = call @softplus_25() : () -> tensor<bf16> loc(#loc270)
    %433 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %434 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %435 = stablehlo.multiply %433, %434 : tensor<f32> loc(#loc271)
    %436 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %437 = call @_where_26(%436) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc272)
    %438 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %439 = stablehlo.broadcast_in_dim %438, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc273)
    %440 = stablehlo.broadcast_in_dim %439, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc274)
    %441 = call @softplus_27() : () -> tensor<bf16> loc(#loc275)
    %442 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %443 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %444 = stablehlo.multiply %442, %443 : tensor<f32> loc(#loc276)
    %445 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %446 = call @_where_28(%445) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc277)
    %447 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %448 = stablehlo.broadcast_in_dim %447, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc278)
    %449 = stablehlo.broadcast_in_dim %448, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc279)
    %450 = call @softplus_29() : () -> tensor<bf16> loc(#loc280)
    %451 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %452 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %453 = stablehlo.multiply %451, %452 : tensor<f32> loc(#loc281)
    %454 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %455 = call @_where_30(%454) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc282)
    %456 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %457 = stablehlo.broadcast_in_dim %456, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc283)
    %458 = stablehlo.broadcast_in_dim %457, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc284)
    %459 = call @softplus_31() : () -> tensor<bf16> loc(#loc285)
    %460 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %461 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %462 = stablehlo.multiply %460, %461 : tensor<f32> loc(#loc286)
    %463 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %464 = call @_where_32(%463) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc287)
    %465 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %466 = stablehlo.broadcast_in_dim %465, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc288)
    %467 = stablehlo.broadcast_in_dim %466, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc289)
    %468 = call @softplus_33() : () -> tensor<bf16> loc(#loc290)
    %469 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %470 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc)
    %471 = stablehlo.multiply %469, %470 : tensor<f32> loc(#loc291)
    %472 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc)
    %473 = call @_where_34(%472) : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc292)
    %474 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %475 = call @_where_35(%474) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %476 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %477 = call @_where_36(%476) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %479 = call @_where_37(%478) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %480 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %481 = call @_where_38(%480) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %482 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %483 = call @_where_39(%482) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %484 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %485 = call @_where_40(%484) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %486 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %487 = call @_where_41(%486) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %488 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %489 = call @_where_42(%488) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %490 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %491 = call @_where_43(%490) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %492 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %493 = call @_where_44(%492) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %494 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %495 = call @_where_45(%494) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %496 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %497 = call @_where_46(%496) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %498 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %499 = call @_where_47(%498) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %500 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %501 = call @_where_48(%500) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %502 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %503 = call @_where_49(%502) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %504 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %505 = call @_where_50(%504) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %506 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %507 = call @_where_51(%506) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %508 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %509 = call @_where_52(%508) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %510 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %511 = call @_where_53(%510) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %512 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %513 = call @_where_54(%512) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %514 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %515 = call @_where_55(%514) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %516 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %517 = call @_where_56(%516) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %518 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %519 = call @_where_57(%518) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %520 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %521 = call @_where_58(%520) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %522 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %523 = call @_where_59(%522) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %524 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %525 = call @_where_60(%524) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %526 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %527 = call @_where_61(%526) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %528 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %529 = call @_where_62(%528) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %530 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %531 = call @_where_63(%530) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %532 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %533 = call @_where_64(%532) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %534 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %535 = call @_where_65(%534) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %536 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %537 = call @_where_66(%536) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %539 = call @_where_67(%538) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %540 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %541 = call @_where_68(%540) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %542 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %543 = call @_where_69(%542) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %544 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %545 = call @_where_70(%544) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %546 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %547 = call @_where_71(%546) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %548 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %549 = call @_where_72(%548) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %550 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %551 = call @_where_73(%550) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %552 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %553 = call @_where_74(%552) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %554 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %555 = call @_where_75(%554) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %556 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %557 = call @_where_76(%556) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %558 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %559 = call @_where_77(%558) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %560 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %561 = call @_where_78(%560) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %562 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %563 = call @_where_79(%562) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %564 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %565 = call @_where_80(%564) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %566 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %567 = call @_where_81(%566) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %568 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %569 = call @_where_82(%568) : (tensor<f32>) -> tensor<1xf32> loc(#loc293)
    %570 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %571:8 = call @_pad_83(%570) : (tensor<i32>) -> (tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>) loc(#loc294)
    %572 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %573 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc295)
    %574 = stablehlo.broadcast_in_dim %573, dims = [] : (tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc296)
    %575 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %576 = stablehlo.broadcast_in_dim %575, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %577 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %578 = stablehlo.broadcast_in_dim %577, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %579 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %580 = stablehlo.broadcast_in_dim %579, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %581 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %582 = stablehlo.broadcast_in_dim %581, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %583 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %584 = stablehlo.broadcast_in_dim %583, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %585 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %586 = stablehlo.broadcast_in_dim %585, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %587 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %588 = stablehlo.broadcast_in_dim %587, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %589 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %590 = stablehlo.broadcast_in_dim %589, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %591 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %592 = stablehlo.broadcast_in_dim %591, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %593 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %594 = stablehlo.broadcast_in_dim %593, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %595 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %596 = stablehlo.broadcast_in_dim %595, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %597 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %598 = stablehlo.broadcast_in_dim %597, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %599 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %600 = stablehlo.broadcast_in_dim %599, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %601 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %602 = stablehlo.broadcast_in_dim %601, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %603 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %604 = stablehlo.broadcast_in_dim %603, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %605 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %606 = stablehlo.broadcast_in_dim %605, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %607 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %608 = stablehlo.broadcast_in_dim %607, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %609 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %610 = stablehlo.broadcast_in_dim %609, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %611 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %612 = stablehlo.broadcast_in_dim %611, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %613 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %614 = stablehlo.broadcast_in_dim %613, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %615 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %616 = stablehlo.broadcast_in_dim %615, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %617 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %618 = stablehlo.broadcast_in_dim %617, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %619 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %620 = stablehlo.broadcast_in_dim %619, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %621 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %622 = stablehlo.broadcast_in_dim %621, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %623 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %624 = stablehlo.broadcast_in_dim %623, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %625 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %626 = stablehlo.broadcast_in_dim %625, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %627 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %628 = stablehlo.broadcast_in_dim %627, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %629 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %630 = stablehlo.broadcast_in_dim %629, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %631 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %632 = stablehlo.broadcast_in_dim %631, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %633 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %634 = stablehlo.broadcast_in_dim %633, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %635 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %636 = stablehlo.broadcast_in_dim %635, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %637 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %638 = stablehlo.broadcast_in_dim %637, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %639 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %640 = stablehlo.broadcast_in_dim %639, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %641 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %642 = stablehlo.broadcast_in_dim %641, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %643 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %644 = stablehlo.broadcast_in_dim %643, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %645 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %646 = stablehlo.broadcast_in_dim %645, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %647 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %648 = stablehlo.broadcast_in_dim %647, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %649 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %650 = stablehlo.broadcast_in_dim %649, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %651 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %652 = stablehlo.broadcast_in_dim %651, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %653 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %654 = stablehlo.broadcast_in_dim %653, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %655 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %656 = stablehlo.broadcast_in_dim %655, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %657 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %658 = stablehlo.broadcast_in_dim %657, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %659 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %660 = stablehlo.broadcast_in_dim %659, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %661 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %662 = stablehlo.broadcast_in_dim %661, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %663 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %664 = stablehlo.broadcast_in_dim %663, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %665 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %666 = stablehlo.broadcast_in_dim %665, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %667 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %668 = stablehlo.broadcast_in_dim %667, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %669 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc297)
    %670 = stablehlo.broadcast_in_dim %669, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc298)
    %671 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc295)
    %672 = stablehlo.broadcast_in_dim %671, dims = [] : (tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc296)
    %673 = stablehlo.constant dense<0> : tensor<i32> loc(#loc299)
    %674 = stablehlo.broadcast_in_dim %673, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc300)
    %675 = stablehlo.constant dense<0> : tensor<i32> loc(#loc301)
    %676:543 = stablehlo.while(%iterArg = %259, %iterArg_0 = %261#0, %iterArg_1 = %261#1, %iterArg_2 = %306, %iterArg_3 = %307, %iterArg_4 = %308, %iterArg_5 = %309, %iterArg_6 = %253, %iterArg_7 = %329, %iterArg_8 = %471, %iterArg_9 = %462, %iterArg_10 = %453, %iterArg_11 = %444, %iterArg_12 = %435, %iterArg_13 = %426, %iterArg_14 = %417, %iterArg_15 = %408, %iterArg_16 = %399, %iterArg_17 = %390, %iterArg_18 = %381, %iterArg_19 = %372, %iterArg_20 = %363, %iterArg_21 = %354, %iterArg_22 = %345, %iterArg_23 = %336, %iterArg_24 = %263#0, %iterArg_25 = %263#1, %iterArg_26 = %256, %iterArg_27 = %267, %iterArg_28 = %310, %iterArg_29 = %24, %iterArg_30 = %24, %iterArg_31 = %23, %iterArg_32 = %23, %iterArg_33 = %25, %iterArg_34 = %25, %iterArg_35 = %26, %iterArg_36 = %26, %iterArg_37 = %333, %iterArg_38 = %338, %iterArg_39 = %27, %iterArg_40 = %27, %iterArg_41 = %22, %iterArg_42 = %22, %iterArg_43 = %21, %iterArg_44 = %21, %iterArg_45 = %18, %iterArg_46 = %18, %iterArg_47 = %17, %iterArg_48 = %17, %iterArg_49 = %280, %iterArg_50 = %281, %iterArg_51 = %250, %iterArg_52 = %20, %iterArg_53 = %20, %iterArg_54 = %19, %iterArg_55 = %19, %iterArg_56 = %35, %iterArg_57 = %35, %iterArg_58 = %34, %iterArg_59 = %34, %iterArg_60 = %36, %iterArg_61 = %36, %iterArg_62 = %37, %iterArg_63 = %37, %iterArg_64 = %342, %iterArg_65 = %347, %iterArg_66 = %38, %iterArg_67 = %38, %iterArg_68 = %33, %iterArg_69 = %33, %iterArg_70 = %32, %iterArg_71 = %32, %iterArg_72 = %29, %iterArg_73 = %29, %iterArg_74 = %28, %iterArg_75 = %28, %iterArg_76 = %31, %iterArg_77 = %31, %iterArg_78 = %30, %iterArg_79 = %30, %iterArg_80 = %112, %iterArg_81 = %112, %iterArg_82 = %111, %iterArg_83 = %111, %iterArg_84 = %113, %iterArg_85 = %113, %iterArg_86 = %114, %iterArg_87 = %114, %iterArg_88 = %351, %iterArg_89 = %356, %iterArg_90 = %115, %iterArg_91 = %115, %iterArg_92 = %110, %iterArg_93 = %110, %iterArg_94 = %109, %iterArg_95 = %109, %iterArg_96 = %106, %iterArg_97 = %106, %iterArg_98 = %105, %iterArg_99 = %105, %iterArg_100 = %108, %iterArg_101 = %108, %iterArg_102 = %107, %iterArg_103 = %107, %iterArg_104 = %123, %iterArg_105 = %123, %iterArg_106 = %122, %iterArg_107 = %122, %iterArg_108 = %124, %iterArg_109 = %124, %iterArg_110 = %125, %iterArg_111 = %125, %iterArg_112 = %360, %iterArg_113 = %365, %iterArg_114 = %126, %iterArg_115 = %126, %iterArg_116 = %121, %iterArg_117 = %121, %iterArg_118 = %120, %iterArg_119 = %120, %iterArg_120 = %117, %iterArg_121 = %117, %iterArg_122 = %116, %iterArg_123 = %116, %iterArg_124 = %119, %iterArg_125 = %119, %iterArg_126 = %118, %iterArg_127 = %118, %iterArg_128 = %134, %iterArg_129 = %134, %iterArg_130 = %133, %iterArg_131 = %133, %iterArg_132 = %135, %iterArg_133 = %135, %iterArg_134 = %136, %iterArg_135 = %136, %iterArg_136 = %369, %iterArg_137 = %374, %iterArg_138 = %137, %iterArg_139 = %137, %iterArg_140 = %132, %iterArg_141 = %132, %iterArg_142 = %131, %iterArg_143 = %131, %iterArg_144 = %128, %iterArg_145 = %128, %iterArg_146 = %127, %iterArg_147 = %127, %iterArg_148 = %130, %iterArg_149 = %130, %iterArg_150 = %129, %iterArg_151 = %129, %iterArg_152 = %145, %iterArg_153 = %145, %iterArg_154 = %144, %iterArg_155 = %144, %iterArg_156 = %146, %iterArg_157 = %146, %iterArg_158 = %147, %iterArg_159 = %147, %iterArg_160 = %378, %iterArg_161 = %383, %iterArg_162 = %148, %iterArg_163 = %148, %iterArg_164 = %143, %iterArg_165 = %143, %iterArg_166 = %142, %iterArg_167 = %142, %iterArg_168 = %139, %iterArg_169 = %139, %iterArg_170 = %138, %iterArg_171 = %138, %iterArg_172 = %141, %iterArg_173 = %141, %iterArg_174 = %140, %iterArg_175 = %140, %iterArg_176 = %156, %iterArg_177 = %156, %iterArg_178 = %155, %iterArg_179 = %155, %iterArg_180 = %157, %iterArg_181 = %157, %iterArg_182 = %158, %iterArg_183 = %158, %iterArg_184 = %387, %iterArg_185 = %392, %iterArg_186 = %159, %iterArg_187 = %159, %iterArg_188 = %154, %iterArg_189 = %154, %iterArg_190 = %153, %iterArg_191 = %153, %iterArg_192 = %150, %iterArg_193 = %150, %iterArg_194 = %149, %iterArg_195 = %149, %iterArg_196 = %152, %iterArg_197 = %152, %iterArg_198 = %151, %iterArg_199 = %151, %iterArg_200 = %167, %iterArg_201 = %167, %iterArg_202 = %166, %iterArg_203 = %166, %iterArg_204 = %168, %iterArg_205 = %168, %iterArg_206 = %169, %iterArg_207 = %169, %iterArg_208 = %396, %iterArg_209 = %401, %iterArg_210 = %170, %iterArg_211 = %170, %iterArg_212 = %165, %iterArg_213 = %165, %iterArg_214 = %164, %iterArg_215 = %164, %iterArg_216 = %161, %iterArg_217 = %161, %iterArg_218 = %160, %iterArg_219 = %160, %iterArg_220 = %163, %iterArg_221 = %163, %iterArg_222 = %162, %iterArg_223 = %162, %iterArg_224 = %178, %iterArg_225 = %178, %iterArg_226 = %177, %iterArg_227 = %177, %iterArg_228 = %179, %iterArg_229 = %179, %iterArg_230 = %180, %iterArg_231 = %180, %iterArg_232 = %405, %iterArg_233 = %410, %iterArg_234 = %181, %iterArg_235 = %181, %iterArg_236 = %176, %iterArg_237 = %176, %iterArg_238 = %175, %iterArg_239 = %175, %iterArg_240 = %172, %iterArg_241 = %172, %iterArg_242 = %171, %iterArg_243 = %171, %iterArg_244 = %174, %iterArg_245 = %174, %iterArg_246 = %173, %iterArg_247 = %173, %iterArg_248 = %189, %iterArg_249 = %189, %iterArg_250 = %188, %iterArg_251 = %188, %iterArg_252 = %190, %iterArg_253 = %190, %iterArg_254 = %191, %iterArg_255 = %191, %iterArg_256 = %414, %iterArg_257 = %419, %iterArg_258 = %192, %iterArg_259 = %192, %iterArg_260 = %187, %iterArg_261 = %187, %iterArg_262 = %186, %iterArg_263 = %186, %iterArg_264 = %183, %iterArg_265 = %183, %iterArg_266 = %182, %iterArg_267 = %182, %iterArg_268 = %185, %iterArg_269 = %185, %iterArg_270 = %184, %iterArg_271 = %184, %iterArg_272 = %46, %iterArg_273 = %46, %iterArg_274 = %45, %iterArg_275 = %45, %iterArg_276 = %47, %iterArg_277 = %47, %iterArg_278 = %48, %iterArg_279 = %48, %iterArg_280 = %423, %iterArg_281 = %428, %iterArg_282 = %49, %iterArg_283 = %49, %iterArg_284 = %44, %iterArg_285 = %44, %iterArg_286 = %43, %iterArg_287 = %43, %iterArg_288 = %40, %iterArg_289 = %40, %iterArg_290 = %39, %iterArg_291 = %39, %iterArg_292 = %42, %iterArg_293 = %42, %iterArg_294 = %41, %iterArg_295 = %41, %iterArg_296 = %57, %iterArg_297 = %57, %iterArg_298 = %56, %iterArg_299 = %56, %iterArg_300 = %58, %iterArg_301 = %58, %iterArg_302 = %59, %iterArg_303 = %59, %iterArg_304 = %432, %iterArg_305 = %437, %iterArg_306 = %60, %iterArg_307 = %60, %iterArg_308 = %55, %iterArg_309 = %55, %iterArg_310 = %54, %iterArg_311 = %54, %iterArg_312 = %51, %iterArg_313 = %51, %iterArg_314 = %50, %iterArg_315 = %50, %iterArg_316 = %53, %iterArg_317 = %53, %iterArg_318 = %52, %iterArg_319 = %52, %iterArg_320 = %68, %iterArg_321 = %68, %iterArg_322 = %67, %iterArg_323 = %67, %iterArg_324 = %69, %iterArg_325 = %69, %iterArg_326 = %70, %iterArg_327 = %70, %iterArg_328 = %441, %iterArg_329 = %446, %iterArg_330 = %71, %iterArg_331 = %71, %iterArg_332 = %66, %iterArg_333 = %66, %iterArg_334 = %65, %iterArg_335 = %65, %iterArg_336 = %62, %iterArg_337 = %62, %iterArg_338 = %61, %iterArg_339 = %61, %iterArg_340 = %64, %iterArg_341 = %64, %iterArg_342 = %63, %iterArg_343 = %63, %iterArg_344 = %79, %iterArg_345 = %79, %iterArg_346 = %78, %iterArg_347 = %78, %iterArg_348 = %80, %iterArg_349 = %80, %iterArg_350 = %81, %iterArg_351 = %81, %iterArg_352 = %450, %iterArg_353 = %455, %iterArg_354 = %82, %iterArg_355 = %82, %iterArg_356 = %77, %iterArg_357 = %77, %iterArg_358 = %76, %iterArg_359 = %76, %iterArg_360 = %73, %iterArg_361 = %73, %iterArg_362 = %72, %iterArg_363 = %72, %iterArg_364 = %75, %iterArg_365 = %75, %iterArg_366 = %74, %iterArg_367 = %74, %iterArg_368 = %90, %iterArg_369 = %90, %iterArg_370 = %89, %iterArg_371 = %89, %iterArg_372 = %91, %iterArg_373 = %91, %iterArg_374 = %92, %iterArg_375 = %92, %iterArg_376 = %459, %iterArg_377 = %464, %iterArg_378 = %93, %iterArg_379 = %93, %iterArg_380 = %88, %iterArg_381 = %88, %iterArg_382 = %87, %iterArg_383 = %87, %iterArg_384 = %84, %iterArg_385 = %84, %iterArg_386 = %83, %iterArg_387 = %83, %iterArg_388 = %86, %iterArg_389 = %86, %iterArg_390 = %85, %iterArg_391 = %85, %iterArg_392 = %101, %iterArg_393 = %101, %iterArg_394 = %100, %iterArg_395 = %100, %iterArg_396 = %102, %iterArg_397 = %102, %iterArg_398 = %103, %iterArg_399 = %103, %iterArg_400 = %468, %iterArg_401 = %473, %iterArg_402 = %104, %iterArg_403 = %104, %iterArg_404 = %99, %iterArg_405 = %99, %iterArg_406 = %98, %iterArg_407 = %98, %iterArg_408 = %95, %iterArg_409 = %95, %iterArg_410 = %94, %iterArg_411 = %94, %iterArg_412 = %97, %iterArg_413 = %97, %iterArg_414 = %96, %iterArg_415 = %96, %iterArg_416 = %571#0, %iterArg_417 = %571#1, %iterArg_418 = %571#2, %iterArg_419 = %571#3, %iterArg_420 = %571#4, %iterArg_421 = %571#5, %iterArg_422 = %571#6, %iterArg_423 = %571#7, %iterArg_424 = %475, %iterArg_425 = %332, %iterArg_426 = %477, %iterArg_427 = %479, %iterArg_428 = %481, %iterArg_429 = %341, %iterArg_430 = %483, %iterArg_431 = %485, %iterArg_432 = %487, %iterArg_433 = %422, %iterArg_434 = %489, %iterArg_435 = %491, %iterArg_436 = %493, %iterArg_437 = %431, %iterArg_438 = %495, %iterArg_439 = %497, %iterArg_440 = %499, %iterArg_441 = %440, %iterArg_442 = %501, %iterArg_443 = %503, %iterArg_444 = %505, %iterArg_445 = %449, %iterArg_446 = %507, %iterArg_447 = %509, %iterArg_448 = %511, %iterArg_449 = %458, %iterArg_450 = %513, %iterArg_451 = %515, %iterArg_452 = %517, %iterArg_453 = %467, %iterArg_454 = %519, %iterArg_455 = %521, %iterArg_456 = %523, %iterArg_457 = %350, %iterArg_458 = %525, %iterArg_459 = %527, %iterArg_460 = %529, %iterArg_461 = %359, %iterArg_462 = %531, %iterArg_463 = %533, %iterArg_464 = %535, %iterArg_465 = %368, %iterArg_466 = %537, %iterArg_467 = %539, %iterArg_468 = %541, %iterArg_469 = %377, %iterArg_470 = %543, %iterArg_471 = %545, %iterArg_472 = %547, %iterArg_473 = %386, %iterArg_474 = %549, %iterArg_475 = %551, %iterArg_476 = %553, %iterArg_477 = %395, %iterArg_478 = %555, %iterArg_479 = %557, %iterArg_480 = %559, %iterArg_481 = %404, %iterArg_482 = %561, %iterArg_483 = %563, %iterArg_484 = %565, %iterArg_485 = %413, %iterArg_486 = %567, %iterArg_487 = %569, %iterArg_488 = %675, %iterArg_489 = %258, %iterArg_490 = %572, %iterArg_491 = %574, %iterArg_492 = %576, %iterArg_493 = %578, %iterArg_494 = %580, %iterArg_495 = %582, %iterArg_496 = %584, %iterArg_497 = %586, %iterArg_498 = %588, %iterArg_499 = %590, %iterArg_500 = %592, %iterArg_501 = %594, %iterArg_502 = %596, %iterArg_503 = %598, %iterArg_504 = %600, %iterArg_505 = %602, %iterArg_506 = %604, %iterArg_507 = %606, %iterArg_508 = %608, %iterArg_509 = %610, %iterArg_510 = %612, %iterArg_511 = %614, %iterArg_512 = %616, %iterArg_513 = %618, %iterArg_514 = %620, %iterArg_515 = %622, %iterArg_516 = %624, %iterArg_517 = %626, %iterArg_518 = %628, %iterArg_519 = %630, %iterArg_520 = %632, %iterArg_521 = %634, %iterArg_522 = %636, %iterArg_523 = %638, %iterArg_524 = %640, %iterArg_525 = %642, %iterArg_526 = %644, %iterArg_527 = %646, %iterArg_528 = %648, %iterArg_529 = %650, %iterArg_530 = %652, %iterArg_531 = %654, %iterArg_532 = %656, %iterArg_533 = %658, %iterArg_534 = %660, %iterArg_535 = %662, %iterArg_536 = %664, %iterArg_537 = %666, %iterArg_538 = %668, %iterArg_539 = %670, %iterArg_540 = %672, %iterArg_541 = %674) : tensor<1xi32>, tensor<i32>, tensor<i1>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1x1x2048x2048xbf16>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<i1>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1x2048x2048xi1>, tensor<1xi32>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1x2048xi32>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<i32>, tensor<1x1x2048x2048xbf16>, tensor<i32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1x1x2048x2048xbf16>, tensor<1xi32>
     cond {
      %29929 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %29930 = stablehlo.compare  LT, %iterArg_488, %29929,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc303)
      stablehlo.return %29930 : tensor<i1> loc(#loc302)
    } do {
      %29929 = stablehlo.broadcast_in_dim %iterArg_490, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc304)
      %29930 = stablehlo.subtract %29929, %iterArg : tensor<1xi32> loc(#loc304)
      %29931 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %29932 = stablehlo.broadcast_in_dim %29931, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc305)
      %29933 = stablehlo.maximum %29930, %29932 : tensor<1xi32> loc(#loc305)
      %29934 = func.call @remainder_84(%29933, %iterArg_0, %iterArg_1) : (tensor<1xi32>, tensor<i32>, tensor<i1>) -> tensor<1xi32> loc(#loc306)
      %29935 = arith.select %ptrue, %29934, %29934 : tensor<1xi32>
      %29936 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %29937 = stablehlo.broadcast_in_dim %29936, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc308)
      %29938 = stablehlo.compare  LT, %29935, %29937,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc308)
      %29939 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %29940 = stablehlo.broadcast_in_dim %29939, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc309)
      %29941 = stablehlo.add %29935, %29940 : tensor<1xi32> loc(#loc309)
      %29942 = stablehlo.select %29938, %29941, %29935 : tensor<1xi1>, tensor<1xi32> loc(#loc310)
      %29943 = stablehlo.broadcast_in_dim %29942, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32> loc(#loc311)
      %29944 = stablehlo.concatenate %29943, %iterArg_2, %iterArg_3, %iterArg_4, %iterArg_5, dim = 1 : (tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<1x5xi32> loc(#loc312)
      %29945 = "stablehlo.gather"(%iterArg_6, %29944) {dimension_numbers = #stablehlo.gather<offset_dims = [1, 2, 3, 4, 5], start_index_map = [0, 1, 2, 3, 4], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>, tensor<1x5xi32>) -> tensor<1x1x1x1x2048x2048xbf16> loc(#loc313)
      %29946 = stablehlo.reshape %29945 : (tensor<1x1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc314)
      %29947 = arith.select %ptrue, %29946, %29946 : tensor<1x1x1x2048x2048xbf16>
      %29948 = arith.select %ptrue, %29947, %29947 : tensor<1x1x1x2048x2048xbf16>
      %29949 = stablehlo.minimum %29948, %iterArg_7 : tensor<1x1x1x2048x2048xbf16> loc(#loc317)
      %29950 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc318)
      %29951 = stablehlo.broadcast_in_dim %iterArg_8, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc319)
      %29952 = stablehlo.compare  GE, %29950, %29951,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc319)
      %29953 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc320)
      %29954 = stablehlo.broadcast_in_dim %iterArg_9, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc321)
      %29955 = stablehlo.compare  GE, %29953, %29954,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc321)
      %29956 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc322)
      %29957 = stablehlo.broadcast_in_dim %iterArg_10, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc323)
      %29958 = stablehlo.compare  GE, %29956, %29957,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc323)
      %29959 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc324)
      %29960 = stablehlo.broadcast_in_dim %iterArg_11, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc325)
      %29961 = stablehlo.compare  GE, %29959, %29960,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc325)
      %29962 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc326)
      %29963 = stablehlo.broadcast_in_dim %iterArg_12, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc327)
      %29964 = stablehlo.compare  GE, %29962, %29963,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc327)
      %29965 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc328)
      %29966 = stablehlo.broadcast_in_dim %iterArg_13, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc329)
      %29967 = stablehlo.compare  GE, %29965, %29966,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc329)
      %29968 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc330)
      %29969 = stablehlo.broadcast_in_dim %iterArg_14, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc331)
      %29970 = stablehlo.compare  GE, %29968, %29969,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc331)
      %29971 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc332)
      %29972 = stablehlo.broadcast_in_dim %iterArg_15, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc333)
      %29973 = stablehlo.compare  GE, %29971, %29972,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc333)
      %29974 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc334)
      %29975 = stablehlo.broadcast_in_dim %iterArg_16, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc335)
      %29976 = stablehlo.compare  GE, %29974, %29975,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc335)
      %29977 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc336)
      %29978 = stablehlo.broadcast_in_dim %iterArg_17, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc337)
      %29979 = stablehlo.compare  GE, %29977, %29978,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc337)
      %29980 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc338)
      %29981 = stablehlo.broadcast_in_dim %iterArg_18, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc339)
      %29982 = stablehlo.compare  GE, %29980, %29981,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc339)
      %29983 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc340)
      %29984 = stablehlo.broadcast_in_dim %iterArg_19, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc341)
      %29985 = stablehlo.compare  GE, %29983, %29984,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc341)
      %29986 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc342)
      %29987 = stablehlo.broadcast_in_dim %iterArg_20, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc343)
      %29988 = stablehlo.compare  GE, %29986, %29987,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc343)
      %29989 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc344)
      %29990 = stablehlo.broadcast_in_dim %iterArg_21, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc345)
      %29991 = stablehlo.compare  GE, %29989, %29990,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc345)
      %29992 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc346)
      %29993 = stablehlo.broadcast_in_dim %iterArg_22, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc347)
      %29994 = stablehlo.compare  GE, %29992, %29993,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc347)
      %29995 = stablehlo.convert %29949 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc348)
      %29996 = stablehlo.broadcast_in_dim %iterArg_23, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc349)
      %29997 = stablehlo.compare  GE, %29995, %29996,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc349)
      %29998 = func.call @remainder_85(%iterArg_490, %iterArg_24, %iterArg_25) : (tensor<i32>, tensor<i32>, tensor<i1>) -> tensor<i32> loc(#loc350)
      %29999 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %30000 = stablehlo.compare  LT, %29998, %29999,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc351)
      %30001 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %30002 = stablehlo.add %29998, %30001 : tensor<i32> loc(#loc352)
      %30003 = stablehlo.select %30000, %30002, %29998 : tensor<i1>, tensor<i32> loc(#loc353)
      %30004 = stablehlo.broadcast_in_dim %30003, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc354)
      %30005 = "stablehlo.gather"(%iterArg_26, %30004) {dimension_numbers = #stablehlo.gather<offset_dims = [0, 1, 2, 3], collapsed_slice_dims = [0], start_index_map = [0]>, indices_are_sorted = true, slice_sizes = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>, tensor<1xi32>) -> tensor<1x1x2048x2048xbf16> loc(#loc355)
      %30006 = func.call @_where_86(%30005, %iterArg_489, %iterArg_27) : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xbf16> loc(#loc356)
      %30007 = arith.select %ptrue, %30006, %30006 : tensor<1x1x2048x2048xbf16>
      %30008 = stablehlo.convert %30007 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc358)
      %30009 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc359)
      %30010 = stablehlo.reduce(%30008 init: %30009) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc359)
        stablehlo.return %33956 : tensor<f32> loc(#loc359)
      } loc(#loc359)
      %30011 = stablehlo.broadcast_in_dim %30010, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc360)
      %30012 = stablehlo.convert %30011 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc361)
      %30013 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30014 = stablehlo.broadcast_in_dim %30013, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc362)
      %30015 = stablehlo.divide %30012, %30014 : tensor<1x1x2048x1xbf16> loc(#loc362)
      %30016 = stablehlo.broadcast_in_dim %30015, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc363)
      %30017 = stablehlo.subtract %30007, %30016 : tensor<1x1x2048x2048xbf16> loc(#loc363)
      %30018 = stablehlo.broadcast_in_dim %30015, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc363)
      %30019 = stablehlo.subtract %30007, %30018 : tensor<1x1x2048x2048xbf16> loc(#loc363)
      %30020 = stablehlo.multiply %30019, %30019 : tensor<1x1x2048x2048xbf16> loc(#loc364)
      %30021 = stablehlo.convert %30020 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc358)
      %30022 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc359)
      %30023 = stablehlo.reduce(%30021 init: %30022) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc359)
        stablehlo.return %33956 : tensor<f32> loc(#loc359)
      } loc(#loc359)
      %30024 = stablehlo.broadcast_in_dim %30023, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc360)
      %30025 = stablehlo.convert %30024 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc361)
      %30026 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30027 = stablehlo.broadcast_in_dim %30026, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc362)
      %30028 = stablehlo.divide %30025, %30027 : tensor<1x1x2048x1xbf16> loc(#loc362)
      %30029 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30030 = stablehlo.broadcast_in_dim %30029, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc365)
      %30031 = stablehlo.add %30028, %30030 : tensor<1x1x2048x1xbf16> loc(#loc365)
      %30032 = stablehlo.rsqrt %30031 : tensor<1x1x2048x1xbf16> loc(#loc366)
      %30033 = stablehlo.broadcast_in_dim %30032, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc367)
      %30034 = stablehlo.multiply %30017, %30033 : tensor<1x1x2048x2048xbf16> loc(#loc367)
      %30035 = stablehlo.broadcast_in_dim %iterArg_490, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc368)
      %30036 = stablehlo.compare  LE, %iterArg_28, %30035,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc368)
      %30037 = stablehlo.broadcast_in_dim %iterArg_490, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc369)
      %30038 = stablehlo.subtract %30037, %iterArg_28 : tensor<1xi32> loc(#loc369)
      %30039 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %30040 = stablehlo.broadcast_in_dim %30039, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc370)
      %30041 = stablehlo.compare  LT, %30038, %30040,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc370)
      %30042 = stablehlo.and %30036, %30041 : tensor<1xi1> loc(#loc371)
      %30043 = func.call @_where_87(%30042, %iterArg_29, %iterArg_30) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30044 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30045 = stablehlo.broadcast_in_dim %30044, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc373)
      %30046 = stablehlo.add %30045, %30043 : tensor<1x2048xbf16> loc(#loc373)
      %30047 = stablehlo.broadcast_in_dim %30046, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc374)
      %30048 = stablehlo.broadcast_in_dim %30047, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc375)
      %30049 = stablehlo.multiply %30034, %30048 : tensor<1x1x2048x2048xbf16> loc(#loc375)
      %30050 = func.call @_where_88(%30042, %iterArg_31, %iterArg_32) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30051 = stablehlo.broadcast_in_dim %30050, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc376)
      %30052 = stablehlo.broadcast_in_dim %30051, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc377)
      %30053 = stablehlo.add %30049, %30052 : tensor<1x1x2048x2048xbf16> loc(#loc377)
      %30054 = func.call @_where_89(%30042, %iterArg_33, %iterArg_34) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %30055 = func.call @_einsum_90(%30053, %30054) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc378)
      %30056 = "stablehlo.slice"(%30055) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc379)
      %30057 = stablehlo.reshape %30056 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc380)
      %30058 = arith.select %ptrue, %30057, %30057 : tensor<1x1x2048x32x64xbf16>
      %30059 = func.call @_where_91(%30042, %iterArg_35, %iterArg_36) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %30060 = func.call @softplus_92(%30059, %iterArg_37) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc382)
      %30061 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %30062 = stablehlo.broadcast_in_dim %30061, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc383)
      %30063 = stablehlo.multiply %30062, %30060 : tensor<1x64xbf16> loc(#loc383)
      %30064 = stablehlo.broadcast_in_dim %30063, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc384)
      %30065 = stablehlo.broadcast_in_dim %30064, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc385)
      %30066 = stablehlo.multiply %30058, %30065 : tensor<1x1x2048x32x64xbf16> loc(#loc385)
      %30067 = "stablehlo.slice"(%30055) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc386)
      %30068 = stablehlo.reshape %30067 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc380)
      %30069 = arith.select %ptrue, %30068, %30068 : tensor<1x1x2048x32x64xbf16>
      %30070 = func.call @_einsum_93(%30066, %30069) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc387)
      %30071 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30072 = stablehlo.broadcast_in_dim %30071, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc388)
      %30073 = stablehlo.divide %30070, %30072 : tensor<1x1x32x2048x2048xbf16> loc(#loc388)
      %30074 = stablehlo.tanh %30073 : tensor<1x1x32x2048x2048xbf16> loc(#loc389)
      %30075 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30076 = stablehlo.broadcast_in_dim %30075, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc390)
      %30077 = stablehlo.multiply %30076, %30074 : tensor<1x1x32x2048x2048xbf16> loc(#loc390)
      %30078 = stablehlo.convert %30077 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc391)
      %30079 = func.call @_where_94(%29997, %30078, %iterArg_38) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc392)
      %30080 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc393)
      %30081 = stablehlo.reduce(%30079 init: %30080) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc393)
        stablehlo.return %33956 : tensor<f32> loc(#loc393)
      } loc(#loc393)
      %30082 = stablehlo.broadcast_in_dim %30081, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc394)
      %30083 = stablehlo.broadcast_in_dim %30082, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc395)
      %30084 = stablehlo.subtract %30079, %30083 : tensor<1x1x32x2048x2048xf32> loc(#loc395)
      %30085 = stablehlo.exponential %30084 : tensor<1x1x32x2048x2048xf32> loc(#loc396)
      %30086 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc397)
      %30087 = stablehlo.reduce(%30085 init: %30086) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc397)
        stablehlo.return %33956 : tensor<f32> loc(#loc397)
      } loc(#loc397)
      %30088 = stablehlo.broadcast_in_dim %30087, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc394)
      %30089 = stablehlo.broadcast_in_dim %30088, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc398)
      %30090 = stablehlo.divide %30085, %30089 : tensor<1x1x32x2048x2048xf32> loc(#loc398)
      %30091 = stablehlo.convert %30090 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc399)
      %30092 = "stablehlo.slice"(%30055) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc400)
      %30093 = stablehlo.reshape %30092 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc380)
      %30094 = arith.select %ptrue, %30093, %30093 : tensor<1x1x2048x32x64xbf16>
      %30095 = func.call @_einsum_95(%30091, %30094) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc401)
      %30096 = arith.select %ptrue, %30095, %30095 : tensor<1x1x2048x32x64xbf16>
      %30097 = func.call @_where_96(%30042, %iterArg_39, %iterArg_40) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %30098 = func.call @_einsum_97(%30096, %30097) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc402)
      %30099 = arith.select %ptrue, %30098, %30098 : tensor<1x1x2048x2048xbf16>
      %30100 = stablehlo.add %30099, %30007 : tensor<1x1x2048x2048xbf16> loc(#loc404)
      %30101 = stablehlo.convert %30100 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc405)
      %30102 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc406)
      %30103 = stablehlo.reduce(%30101 init: %30102) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc406)
        stablehlo.return %33956 : tensor<f32> loc(#loc406)
      } loc(#loc406)
      %30104 = stablehlo.broadcast_in_dim %30103, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc407)
      %30105 = stablehlo.convert %30104 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc408)
      %30106 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30107 = stablehlo.broadcast_in_dim %30106, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc409)
      %30108 = stablehlo.divide %30105, %30107 : tensor<1x1x2048x1xbf16> loc(#loc409)
      %30109 = stablehlo.broadcast_in_dim %30108, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc410)
      %30110 = stablehlo.subtract %30100, %30109 : tensor<1x1x2048x2048xbf16> loc(#loc410)
      %30111 = stablehlo.broadcast_in_dim %30108, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc410)
      %30112 = stablehlo.subtract %30100, %30111 : tensor<1x1x2048x2048xbf16> loc(#loc410)
      %30113 = stablehlo.multiply %30112, %30112 : tensor<1x1x2048x2048xbf16> loc(#loc411)
      %30114 = stablehlo.convert %30113 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc405)
      %30115 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc406)
      %30116 = stablehlo.reduce(%30114 init: %30115) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc406)
        stablehlo.return %33956 : tensor<f32> loc(#loc406)
      } loc(#loc406)
      %30117 = stablehlo.broadcast_in_dim %30116, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc407)
      %30118 = stablehlo.convert %30117 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc408)
      %30119 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30120 = stablehlo.broadcast_in_dim %30119, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc409)
      %30121 = stablehlo.divide %30118, %30120 : tensor<1x1x2048x1xbf16> loc(#loc409)
      %30122 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30123 = stablehlo.broadcast_in_dim %30122, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc412)
      %30124 = stablehlo.add %30121, %30123 : tensor<1x1x2048x1xbf16> loc(#loc412)
      %30125 = stablehlo.rsqrt %30124 : tensor<1x1x2048x1xbf16> loc(#loc413)
      %30126 = stablehlo.broadcast_in_dim %30125, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc414)
      %30127 = stablehlo.multiply %30110, %30126 : tensor<1x1x2048x2048xbf16> loc(#loc414)
      %30128 = func.call @_where_98(%30042, %iterArg_41, %iterArg_42) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30129 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30130 = stablehlo.broadcast_in_dim %30129, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc415)
      %30131 = stablehlo.add %30130, %30128 : tensor<1x2048xbf16> loc(#loc415)
      %30132 = stablehlo.broadcast_in_dim %30131, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc416)
      %30133 = stablehlo.broadcast_in_dim %30132, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc417)
      %30134 = stablehlo.multiply %30127, %30133 : tensor<1x1x2048x2048xbf16> loc(#loc417)
      %30135 = func.call @_where_99(%30042, %iterArg_43, %iterArg_44) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30136 = stablehlo.broadcast_in_dim %30135, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc418)
      %30137 = stablehlo.broadcast_in_dim %30136, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc419)
      %30138 = stablehlo.add %30134, %30137 : tensor<1x1x2048x2048xbf16> loc(#loc419)
      %30139 = func.call @_where_100(%30042, %iterArg_45, %iterArg_46) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %30140 = func.call @_einsum_101(%30138, %30139) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc420)
      %30141 = arith.select %ptrue, %30140, %30140 : tensor<1x1x2048x8192xbf16>
      %30142 = func.call @_where_102(%30042, %iterArg_47, %iterArg_48) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %30143 = stablehlo.broadcast_in_dim %30142, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc422)
      %30144 = stablehlo.broadcast_in_dim %30143, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc423)
      %30145 = stablehlo.add %30141, %30144 : tensor<1x1x2048x8192xbf16> loc(#loc423)
      %30146 = stablehlo.multiply %30145, %30145 : tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30147 = stablehlo.multiply %30145, %30146 : tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30148 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %30149 = stablehlo.broadcast_in_dim %30148, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30150 = stablehlo.multiply %30149, %30147 : tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30151 = stablehlo.add %30145, %30150 : tensor<1x1x2048x8192xbf16> loc(#loc425)
      %30152 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %30153 = stablehlo.broadcast_in_dim %30152, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30154 = stablehlo.multiply %30153, %30151 : tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30155 = stablehlo.tanh %30154 : tensor<1x1x2048x8192xbf16> loc(#loc426)
      %30156 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30157 = stablehlo.broadcast_in_dim %30156, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc425)
      %30158 = stablehlo.add %30157, %30155 : tensor<1x1x2048x8192xbf16> loc(#loc425)
      %30159 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %30160 = stablehlo.broadcast_in_dim %30159, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30161 = stablehlo.multiply %30160, %30158 : tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30162 = stablehlo.multiply %30145, %30161 : tensor<1x1x2048x8192xbf16> loc(#loc424)
      %30163 = arith.select %ptrue, %29934, %29934 : tensor<1xi32>
      %30164 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %30165 = stablehlo.broadcast_in_dim %30164, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc308)
      %30166 = stablehlo.compare  LT, %30163, %30165,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc308)
      %30167 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %30168 = stablehlo.broadcast_in_dim %30167, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc309)
      %30169 = stablehlo.add %30163, %30168 : tensor<1xi32> loc(#loc309)
      %30170 = stablehlo.select %30166, %30169, %30163 : tensor<1xi1>, tensor<1xi32> loc(#loc310)
      %30171 = stablehlo.broadcast_in_dim %30170, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32> loc(#loc311)
      %30172 = stablehlo.concatenate %30171, %iterArg_49, %iterArg_50, dim = 1 : (tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<1x3xi32> loc(#loc312)
      %30173 = "stablehlo.gather"(%iterArg_51, %30172) {dimension_numbers = #stablehlo.gather<offset_dims = [1, 2, 3], start_index_map = [0, 1, 2], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = dense<[1, 1, 2048]> : tensor<3xi64>} : (tensor<1x1x2048xi32>, tensor<1x3xi32>) -> tensor<1x1x1x2048xi32> loc(#loc427)
      %30174 = stablehlo.reshape %30173 : (tensor<1x1x1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc314)
      %30175 = arith.select %ptrue, %30174, %30174 : tensor<1x1x2048xi32>
      %30176 = arith.select %ptrue, %30175, %30175 : tensor<1x1x2048xi32>
      %30177 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc430)
      %30178 = stablehlo.convert %30177 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc431)
      %30179 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30180 = stablehlo.broadcast_in_dim %30179, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc432)
      %30181 = stablehlo.subtract %30180, %30178 : tensor<1x1x2048x1xf32> loc(#loc432)
      %30182 = stablehlo.convert %30181 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc433)
      %30183 = stablehlo.broadcast_in_dim %30182, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc434)
      %30184 = stablehlo.multiply %30162, %30183 : tensor<1x1x2048x8192xbf16> loc(#loc434)
      %30185 = func.call @_where_103(%30042, %iterArg_52, %iterArg_53) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %30186 = func.call @_einsum_104(%30184, %30185) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc435)
      %30187 = arith.select %ptrue, %30186, %30186 : tensor<1x1x2048x2048xbf16>
      %30188 = func.call @_where_105(%30042, %iterArg_54, %iterArg_55) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30189 = stablehlo.broadcast_in_dim %30188, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc437)
      %30190 = stablehlo.broadcast_in_dim %30189, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc438)
      %30191 = stablehlo.add %30187, %30190 : tensor<1x1x2048x2048xbf16> loc(#loc438)
      %30192 = stablehlo.convert %30177 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc431)
      %30193 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30194 = stablehlo.broadcast_in_dim %30193, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc432)
      %30195 = stablehlo.subtract %30194, %30192 : tensor<1x1x2048x1xf32> loc(#loc432)
      %30196 = stablehlo.convert %30195 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc439)
      %30197 = stablehlo.broadcast_in_dim %30196, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc440)
      %30198 = stablehlo.multiply %30191, %30197 : tensor<1x1x2048x2048xbf16> loc(#loc440)
      %30199 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30200 = stablehlo.broadcast_in_dim %30199, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc441)
      %30201 = stablehlo.multiply %30198, %30200 : tensor<1x1x2048x2048xbf16> loc(#loc441)
      %30202 = stablehlo.add %30100, %30201 : tensor<1x1x2048x2048xbf16> loc(#loc442)
      %30203 = stablehlo.convert %30202 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc443)
      %30204 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc444)
      %30205 = stablehlo.reduce(%30203 init: %30204) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc444)
        stablehlo.return %33956 : tensor<f32> loc(#loc444)
      } loc(#loc444)
      %30206 = stablehlo.broadcast_in_dim %30205, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc445)
      %30207 = stablehlo.convert %30206 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc446)
      %30208 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30209 = stablehlo.broadcast_in_dim %30208, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc447)
      %30210 = stablehlo.divide %30207, %30209 : tensor<1x1x2048x1xbf16> loc(#loc447)
      %30211 = stablehlo.broadcast_in_dim %30210, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc448)
      %30212 = stablehlo.subtract %30202, %30211 : tensor<1x1x2048x2048xbf16> loc(#loc448)
      %30213 = stablehlo.broadcast_in_dim %30210, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc448)
      %30214 = stablehlo.subtract %30202, %30213 : tensor<1x1x2048x2048xbf16> loc(#loc448)
      %30215 = stablehlo.multiply %30214, %30214 : tensor<1x1x2048x2048xbf16> loc(#loc449)
      %30216 = stablehlo.convert %30215 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc443)
      %30217 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc444)
      %30218 = stablehlo.reduce(%30216 init: %30217) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc444)
        stablehlo.return %33956 : tensor<f32> loc(#loc444)
      } loc(#loc444)
      %30219 = stablehlo.broadcast_in_dim %30218, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc445)
      %30220 = stablehlo.convert %30219 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc446)
      %30221 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30222 = stablehlo.broadcast_in_dim %30221, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc447)
      %30223 = stablehlo.divide %30220, %30222 : tensor<1x1x2048x1xbf16> loc(#loc447)
      %30224 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30225 = stablehlo.broadcast_in_dim %30224, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc450)
      %30226 = stablehlo.add %30223, %30225 : tensor<1x1x2048x1xbf16> loc(#loc450)
      %30227 = stablehlo.rsqrt %30226 : tensor<1x1x2048x1xbf16> loc(#loc451)
      %30228 = stablehlo.broadcast_in_dim %30227, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc452)
      %30229 = stablehlo.multiply %30212, %30228 : tensor<1x1x2048x2048xbf16> loc(#loc452)
      %30230 = func.call @_where_106(%30042, %iterArg_56, %iterArg_57) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30231 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30232 = stablehlo.broadcast_in_dim %30231, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc453)
      %30233 = stablehlo.add %30232, %30230 : tensor<1x2048xbf16> loc(#loc453)
      %30234 = stablehlo.broadcast_in_dim %30233, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc454)
      %30235 = stablehlo.broadcast_in_dim %30234, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc455)
      %30236 = stablehlo.multiply %30229, %30235 : tensor<1x1x2048x2048xbf16> loc(#loc455)
      %30237 = func.call @_where_107(%30042, %iterArg_58, %iterArg_59) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30238 = stablehlo.broadcast_in_dim %30237, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc456)
      %30239 = stablehlo.broadcast_in_dim %30238, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc457)
      %30240 = stablehlo.add %30236, %30239 : tensor<1x1x2048x2048xbf16> loc(#loc457)
      %30241 = func.call @_where_108(%30042, %iterArg_60, %iterArg_61) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %30242 = func.call @_einsum_109(%30240, %30241) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc458)
      %30243 = "stablehlo.slice"(%30242) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc459)
      %30244 = stablehlo.reshape %30243 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc460)
      %30245 = arith.select %ptrue, %30244, %30244 : tensor<1x1x2048x32x64xbf16>
      %30246 = func.call @_where_110(%30042, %iterArg_62, %iterArg_63) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %30247 = func.call @softplus_111(%30246, %iterArg_64) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc462)
      %30248 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %30249 = stablehlo.broadcast_in_dim %30248, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc463)
      %30250 = stablehlo.multiply %30249, %30247 : tensor<1x64xbf16> loc(#loc463)
      %30251 = stablehlo.broadcast_in_dim %30250, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc464)
      %30252 = stablehlo.broadcast_in_dim %30251, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc465)
      %30253 = stablehlo.multiply %30245, %30252 : tensor<1x1x2048x32x64xbf16> loc(#loc465)
      %30254 = "stablehlo.slice"(%30242) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc466)
      %30255 = stablehlo.reshape %30254 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc460)
      %30256 = arith.select %ptrue, %30255, %30255 : tensor<1x1x2048x32x64xbf16>
      %30257 = func.call @_einsum_113(%30253, %30256) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc467)
      %30258 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30259 = stablehlo.broadcast_in_dim %30258, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc468)
      %30260 = stablehlo.divide %30257, %30259 : tensor<1x1x32x2048x2048xbf16> loc(#loc468)
      %30261 = stablehlo.tanh %30260 : tensor<1x1x32x2048x2048xbf16> loc(#loc469)
      %30262 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30263 = stablehlo.broadcast_in_dim %30262, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc470)
      %30264 = stablehlo.multiply %30263, %30261 : tensor<1x1x32x2048x2048xbf16> loc(#loc470)
      %30265 = stablehlo.convert %30264 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc471)
      %30266 = func.call @_where_114(%29994, %30265, %iterArg_65) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc472)
      %30267 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc473)
      %30268 = stablehlo.reduce(%30266 init: %30267) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc473)
        stablehlo.return %33956 : tensor<f32> loc(#loc473)
      } loc(#loc473)
      %30269 = stablehlo.broadcast_in_dim %30268, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc474)
      %30270 = stablehlo.broadcast_in_dim %30269, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc475)
      %30271 = stablehlo.subtract %30266, %30270 : tensor<1x1x32x2048x2048xf32> loc(#loc475)
      %30272 = stablehlo.exponential %30271 : tensor<1x1x32x2048x2048xf32> loc(#loc476)
      %30273 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc477)
      %30274 = stablehlo.reduce(%30272 init: %30273) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc477)
        stablehlo.return %33956 : tensor<f32> loc(#loc477)
      } loc(#loc477)
      %30275 = stablehlo.broadcast_in_dim %30274, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc474)
      %30276 = stablehlo.broadcast_in_dim %30275, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc478)
      %30277 = stablehlo.divide %30272, %30276 : tensor<1x1x32x2048x2048xf32> loc(#loc478)
      %30278 = stablehlo.convert %30277 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc479)
      %30279 = "stablehlo.slice"(%30242) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc480)
      %30280 = stablehlo.reshape %30279 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc460)
      %30281 = arith.select %ptrue, %30280, %30280 : tensor<1x1x2048x32x64xbf16>
      %30282 = func.call @_einsum_115(%30278, %30281) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc481)
      %30283 = arith.select %ptrue, %30282, %30282 : tensor<1x1x2048x32x64xbf16>
      %30284 = func.call @_where_116(%30042, %iterArg_66, %iterArg_67) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %30285 = func.call @_einsum_117(%30283, %30284) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc482)
      %30286 = arith.select %ptrue, %30285, %30285 : tensor<1x1x2048x2048xbf16>
      %30287 = stablehlo.add %30286, %30202 : tensor<1x1x2048x2048xbf16> loc(#loc484)
      %30288 = stablehlo.convert %30287 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc485)
      %30289 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc486)
      %30290 = stablehlo.reduce(%30288 init: %30289) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc486)
        stablehlo.return %33956 : tensor<f32> loc(#loc486)
      } loc(#loc486)
      %30291 = stablehlo.broadcast_in_dim %30290, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc487)
      %30292 = stablehlo.convert %30291 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc488)
      %30293 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30294 = stablehlo.broadcast_in_dim %30293, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc489)
      %30295 = stablehlo.divide %30292, %30294 : tensor<1x1x2048x1xbf16> loc(#loc489)
      %30296 = stablehlo.broadcast_in_dim %30295, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc490)
      %30297 = stablehlo.subtract %30287, %30296 : tensor<1x1x2048x2048xbf16> loc(#loc490)
      %30298 = stablehlo.broadcast_in_dim %30295, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc490)
      %30299 = stablehlo.subtract %30287, %30298 : tensor<1x1x2048x2048xbf16> loc(#loc490)
      %30300 = stablehlo.multiply %30299, %30299 : tensor<1x1x2048x2048xbf16> loc(#loc491)
      %30301 = stablehlo.convert %30300 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc485)
      %30302 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc486)
      %30303 = stablehlo.reduce(%30301 init: %30302) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc486)
        stablehlo.return %33956 : tensor<f32> loc(#loc486)
      } loc(#loc486)
      %30304 = stablehlo.broadcast_in_dim %30303, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc487)
      %30305 = stablehlo.convert %30304 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc488)
      %30306 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30307 = stablehlo.broadcast_in_dim %30306, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc489)
      %30308 = stablehlo.divide %30305, %30307 : tensor<1x1x2048x1xbf16> loc(#loc489)
      %30309 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30310 = stablehlo.broadcast_in_dim %30309, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc492)
      %30311 = stablehlo.add %30308, %30310 : tensor<1x1x2048x1xbf16> loc(#loc492)
      %30312 = stablehlo.rsqrt %30311 : tensor<1x1x2048x1xbf16> loc(#loc493)
      %30313 = stablehlo.broadcast_in_dim %30312, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc494)
      %30314 = stablehlo.multiply %30297, %30313 : tensor<1x1x2048x2048xbf16> loc(#loc494)
      %30315 = func.call @_where_118(%30042, %iterArg_68, %iterArg_69) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30316 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30317 = stablehlo.broadcast_in_dim %30316, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc495)
      %30318 = stablehlo.add %30317, %30315 : tensor<1x2048xbf16> loc(#loc495)
      %30319 = stablehlo.broadcast_in_dim %30318, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc496)
      %30320 = stablehlo.broadcast_in_dim %30319, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc497)
      %30321 = stablehlo.multiply %30314, %30320 : tensor<1x1x2048x2048xbf16> loc(#loc497)
      %30322 = func.call @_where_119(%30042, %iterArg_70, %iterArg_71) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30323 = stablehlo.broadcast_in_dim %30322, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc498)
      %30324 = stablehlo.broadcast_in_dim %30323, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc499)
      %30325 = stablehlo.add %30321, %30324 : tensor<1x1x2048x2048xbf16> loc(#loc499)
      %30326 = func.call @_where_120(%30042, %iterArg_72, %iterArg_73) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %30327 = func.call @_einsum_121(%30325, %30326) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc500)
      %30328 = arith.select %ptrue, %30327, %30327 : tensor<1x1x2048x8192xbf16>
      %30329 = func.call @_where_122(%30042, %iterArg_74, %iterArg_75) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %30330 = stablehlo.broadcast_in_dim %30329, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc502)
      %30331 = stablehlo.broadcast_in_dim %30330, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc503)
      %30332 = stablehlo.add %30328, %30331 : tensor<1x1x2048x8192xbf16> loc(#loc503)
      %30333 = stablehlo.multiply %30332, %30332 : tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30334 = stablehlo.multiply %30332, %30333 : tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30335 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %30336 = stablehlo.broadcast_in_dim %30335, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30337 = stablehlo.multiply %30336, %30334 : tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30338 = stablehlo.add %30332, %30337 : tensor<1x1x2048x8192xbf16> loc(#loc505)
      %30339 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %30340 = stablehlo.broadcast_in_dim %30339, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30341 = stablehlo.multiply %30340, %30338 : tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30342 = stablehlo.tanh %30341 : tensor<1x1x2048x8192xbf16> loc(#loc506)
      %30343 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30344 = stablehlo.broadcast_in_dim %30343, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc505)
      %30345 = stablehlo.add %30344, %30342 : tensor<1x1x2048x8192xbf16> loc(#loc505)
      %30346 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %30347 = stablehlo.broadcast_in_dim %30346, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30348 = stablehlo.multiply %30347, %30345 : tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30349 = stablehlo.multiply %30332, %30348 : tensor<1x1x2048x8192xbf16> loc(#loc504)
      %30350 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc507)
      %30351 = stablehlo.convert %30350 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc508)
      %30352 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30353 = stablehlo.broadcast_in_dim %30352, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc509)
      %30354 = stablehlo.subtract %30353, %30351 : tensor<1x1x2048x1xf32> loc(#loc509)
      %30355 = stablehlo.convert %30354 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc510)
      %30356 = stablehlo.broadcast_in_dim %30355, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc511)
      %30357 = stablehlo.multiply %30349, %30356 : tensor<1x1x2048x8192xbf16> loc(#loc511)
      %30358 = func.call @_where_123(%30042, %iterArg_76, %iterArg_77) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %30359 = func.call @_einsum_124(%30357, %30358) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc512)
      %30360 = arith.select %ptrue, %30359, %30359 : tensor<1x1x2048x2048xbf16>
      %30361 = func.call @_where_125(%30042, %iterArg_78, %iterArg_79) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30362 = stablehlo.broadcast_in_dim %30361, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc514)
      %30363 = stablehlo.broadcast_in_dim %30362, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc515)
      %30364 = stablehlo.add %30360, %30363 : tensor<1x1x2048x2048xbf16> loc(#loc515)
      %30365 = stablehlo.convert %30350 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc508)
      %30366 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30367 = stablehlo.broadcast_in_dim %30366, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc509)
      %30368 = stablehlo.subtract %30367, %30365 : tensor<1x1x2048x1xf32> loc(#loc509)
      %30369 = stablehlo.convert %30368 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc516)
      %30370 = stablehlo.broadcast_in_dim %30369, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc517)
      %30371 = stablehlo.multiply %30364, %30370 : tensor<1x1x2048x2048xbf16> loc(#loc517)
      %30372 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30373 = stablehlo.broadcast_in_dim %30372, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc518)
      %30374 = stablehlo.multiply %30371, %30373 : tensor<1x1x2048x2048xbf16> loc(#loc518)
      %30375 = stablehlo.add %30287, %30374 : tensor<1x1x2048x2048xbf16> loc(#loc519)
      %30376 = stablehlo.convert %30375 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc520)
      %30377 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc521)
      %30378 = stablehlo.reduce(%30376 init: %30377) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc521)
        stablehlo.return %33956 : tensor<f32> loc(#loc521)
      } loc(#loc521)
      %30379 = stablehlo.broadcast_in_dim %30378, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc522)
      %30380 = stablehlo.convert %30379 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc523)
      %30381 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30382 = stablehlo.broadcast_in_dim %30381, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc524)
      %30383 = stablehlo.divide %30380, %30382 : tensor<1x1x2048x1xbf16> loc(#loc524)
      %30384 = stablehlo.broadcast_in_dim %30383, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc525)
      %30385 = stablehlo.subtract %30375, %30384 : tensor<1x1x2048x2048xbf16> loc(#loc525)
      %30386 = stablehlo.broadcast_in_dim %30383, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc525)
      %30387 = stablehlo.subtract %30375, %30386 : tensor<1x1x2048x2048xbf16> loc(#loc525)
      %30388 = stablehlo.multiply %30387, %30387 : tensor<1x1x2048x2048xbf16> loc(#loc526)
      %30389 = stablehlo.convert %30388 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc520)
      %30390 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc521)
      %30391 = stablehlo.reduce(%30389 init: %30390) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc521)
        stablehlo.return %33956 : tensor<f32> loc(#loc521)
      } loc(#loc521)
      %30392 = stablehlo.broadcast_in_dim %30391, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc522)
      %30393 = stablehlo.convert %30392 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc523)
      %30394 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30395 = stablehlo.broadcast_in_dim %30394, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc524)
      %30396 = stablehlo.divide %30393, %30395 : tensor<1x1x2048x1xbf16> loc(#loc524)
      %30397 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30398 = stablehlo.broadcast_in_dim %30397, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc527)
      %30399 = stablehlo.add %30396, %30398 : tensor<1x1x2048x1xbf16> loc(#loc527)
      %30400 = stablehlo.rsqrt %30399 : tensor<1x1x2048x1xbf16> loc(#loc528)
      %30401 = stablehlo.broadcast_in_dim %30400, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc529)
      %30402 = stablehlo.multiply %30385, %30401 : tensor<1x1x2048x2048xbf16> loc(#loc529)
      %30403 = func.call @_where_126(%30042, %iterArg_80, %iterArg_81) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30404 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30405 = stablehlo.broadcast_in_dim %30404, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc530)
      %30406 = stablehlo.add %30405, %30403 : tensor<1x2048xbf16> loc(#loc530)
      %30407 = stablehlo.broadcast_in_dim %30406, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc531)
      %30408 = stablehlo.broadcast_in_dim %30407, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc532)
      %30409 = stablehlo.multiply %30402, %30408 : tensor<1x1x2048x2048xbf16> loc(#loc532)
      %30410 = func.call @_where_127(%30042, %iterArg_82, %iterArg_83) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30411 = stablehlo.broadcast_in_dim %30410, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc533)
      %30412 = stablehlo.broadcast_in_dim %30411, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc534)
      %30413 = stablehlo.add %30409, %30412 : tensor<1x1x2048x2048xbf16> loc(#loc534)
      %30414 = func.call @_where_128(%30042, %iterArg_84, %iterArg_85) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %30415 = func.call @_einsum_129(%30413, %30414) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc535)
      %30416 = "stablehlo.slice"(%30415) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc536)
      %30417 = stablehlo.reshape %30416 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc537)
      %30418 = arith.select %ptrue, %30417, %30417 : tensor<1x1x2048x32x64xbf16>
      %30419 = func.call @_where_130(%30042, %iterArg_86, %iterArg_87) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %30420 = func.call @softplus_131(%30419, %iterArg_88) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc539)
      %30421 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %30422 = stablehlo.broadcast_in_dim %30421, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc540)
      %30423 = stablehlo.multiply %30422, %30420 : tensor<1x64xbf16> loc(#loc540)
      %30424 = stablehlo.broadcast_in_dim %30423, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc541)
      %30425 = stablehlo.broadcast_in_dim %30424, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc542)
      %30426 = stablehlo.multiply %30418, %30425 : tensor<1x1x2048x32x64xbf16> loc(#loc542)
      %30427 = "stablehlo.slice"(%30415) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc543)
      %30428 = stablehlo.reshape %30427 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc537)
      %30429 = arith.select %ptrue, %30428, %30428 : tensor<1x1x2048x32x64xbf16>
      %30430 = func.call @_einsum_133(%30426, %30429) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc544)
      %30431 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30432 = stablehlo.broadcast_in_dim %30431, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc545)
      %30433 = stablehlo.divide %30430, %30432 : tensor<1x1x32x2048x2048xbf16> loc(#loc545)
      %30434 = stablehlo.tanh %30433 : tensor<1x1x32x2048x2048xbf16> loc(#loc546)
      %30435 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30436 = stablehlo.broadcast_in_dim %30435, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc547)
      %30437 = stablehlo.multiply %30436, %30434 : tensor<1x1x32x2048x2048xbf16> loc(#loc547)
      %30438 = stablehlo.convert %30437 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc548)
      %30439 = func.call @_where_134(%29991, %30438, %iterArg_89) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc549)
      %30440 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc550)
      %30441 = stablehlo.reduce(%30439 init: %30440) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc550)
        stablehlo.return %33956 : tensor<f32> loc(#loc550)
      } loc(#loc550)
      %30442 = stablehlo.broadcast_in_dim %30441, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc551)
      %30443 = stablehlo.broadcast_in_dim %30442, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc552)
      %30444 = stablehlo.subtract %30439, %30443 : tensor<1x1x32x2048x2048xf32> loc(#loc552)
      %30445 = stablehlo.exponential %30444 : tensor<1x1x32x2048x2048xf32> loc(#loc553)
      %30446 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc554)
      %30447 = stablehlo.reduce(%30445 init: %30446) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc554)
        stablehlo.return %33956 : tensor<f32> loc(#loc554)
      } loc(#loc554)
      %30448 = stablehlo.broadcast_in_dim %30447, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc551)
      %30449 = stablehlo.broadcast_in_dim %30448, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc555)
      %30450 = stablehlo.divide %30445, %30449 : tensor<1x1x32x2048x2048xf32> loc(#loc555)
      %30451 = stablehlo.convert %30450 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc556)
      %30452 = "stablehlo.slice"(%30415) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc557)
      %30453 = stablehlo.reshape %30452 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc537)
      %30454 = arith.select %ptrue, %30453, %30453 : tensor<1x1x2048x32x64xbf16>
      %30455 = func.call @_einsum_135(%30451, %30454) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc558)
      %30456 = arith.select %ptrue, %30455, %30455 : tensor<1x1x2048x32x64xbf16>
      %30457 = func.call @_where_136(%30042, %iterArg_90, %iterArg_91) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %30458 = func.call @_einsum_137(%30456, %30457) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc559)
      %30459 = arith.select %ptrue, %30458, %30458 : tensor<1x1x2048x2048xbf16>
      %30460 = stablehlo.add %30459, %30375 : tensor<1x1x2048x2048xbf16> loc(#loc561)
      %30461 = stablehlo.convert %30460 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc562)
      %30462 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc563)
      %30463 = stablehlo.reduce(%30461 init: %30462) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc563)
        stablehlo.return %33956 : tensor<f32> loc(#loc563)
      } loc(#loc563)
      %30464 = stablehlo.broadcast_in_dim %30463, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc564)
      %30465 = stablehlo.convert %30464 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc565)
      %30466 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30467 = stablehlo.broadcast_in_dim %30466, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc566)
      %30468 = stablehlo.divide %30465, %30467 : tensor<1x1x2048x1xbf16> loc(#loc566)
      %30469 = stablehlo.broadcast_in_dim %30468, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc567)
      %30470 = stablehlo.subtract %30460, %30469 : tensor<1x1x2048x2048xbf16> loc(#loc567)
      %30471 = stablehlo.broadcast_in_dim %30468, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc567)
      %30472 = stablehlo.subtract %30460, %30471 : tensor<1x1x2048x2048xbf16> loc(#loc567)
      %30473 = stablehlo.multiply %30472, %30472 : tensor<1x1x2048x2048xbf16> loc(#loc568)
      %30474 = stablehlo.convert %30473 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc562)
      %30475 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc563)
      %30476 = stablehlo.reduce(%30474 init: %30475) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc563)
        stablehlo.return %33956 : tensor<f32> loc(#loc563)
      } loc(#loc563)
      %30477 = stablehlo.broadcast_in_dim %30476, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc564)
      %30478 = stablehlo.convert %30477 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc565)
      %30479 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30480 = stablehlo.broadcast_in_dim %30479, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc566)
      %30481 = stablehlo.divide %30478, %30480 : tensor<1x1x2048x1xbf16> loc(#loc566)
      %30482 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30483 = stablehlo.broadcast_in_dim %30482, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc569)
      %30484 = stablehlo.add %30481, %30483 : tensor<1x1x2048x1xbf16> loc(#loc569)
      %30485 = stablehlo.rsqrt %30484 : tensor<1x1x2048x1xbf16> loc(#loc570)
      %30486 = stablehlo.broadcast_in_dim %30485, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc571)
      %30487 = stablehlo.multiply %30470, %30486 : tensor<1x1x2048x2048xbf16> loc(#loc571)
      %30488 = func.call @_where_138(%30042, %iterArg_92, %iterArg_93) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30489 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30490 = stablehlo.broadcast_in_dim %30489, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc572)
      %30491 = stablehlo.add %30490, %30488 : tensor<1x2048xbf16> loc(#loc572)
      %30492 = stablehlo.broadcast_in_dim %30491, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc573)
      %30493 = stablehlo.broadcast_in_dim %30492, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc574)
      %30494 = stablehlo.multiply %30487, %30493 : tensor<1x1x2048x2048xbf16> loc(#loc574)
      %30495 = func.call @_where_139(%30042, %iterArg_94, %iterArg_95) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30496 = stablehlo.broadcast_in_dim %30495, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc575)
      %30497 = stablehlo.broadcast_in_dim %30496, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc576)
      %30498 = stablehlo.add %30494, %30497 : tensor<1x1x2048x2048xbf16> loc(#loc576)
      %30499 = func.call @_where_140(%30042, %iterArg_96, %iterArg_97) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %30500 = func.call @_einsum_141(%30498, %30499) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc577)
      %30501 = arith.select %ptrue, %30500, %30500 : tensor<1x1x2048x8192xbf16>
      %30502 = func.call @_where_142(%30042, %iterArg_98, %iterArg_99) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %30503 = stablehlo.broadcast_in_dim %30502, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc579)
      %30504 = stablehlo.broadcast_in_dim %30503, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc580)
      %30505 = stablehlo.add %30501, %30504 : tensor<1x1x2048x8192xbf16> loc(#loc580)
      %30506 = stablehlo.multiply %30505, %30505 : tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30507 = stablehlo.multiply %30505, %30506 : tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30508 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %30509 = stablehlo.broadcast_in_dim %30508, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30510 = stablehlo.multiply %30509, %30507 : tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30511 = stablehlo.add %30505, %30510 : tensor<1x1x2048x8192xbf16> loc(#loc582)
      %30512 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %30513 = stablehlo.broadcast_in_dim %30512, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30514 = stablehlo.multiply %30513, %30511 : tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30515 = stablehlo.tanh %30514 : tensor<1x1x2048x8192xbf16> loc(#loc583)
      %30516 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30517 = stablehlo.broadcast_in_dim %30516, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc582)
      %30518 = stablehlo.add %30517, %30515 : tensor<1x1x2048x8192xbf16> loc(#loc582)
      %30519 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %30520 = stablehlo.broadcast_in_dim %30519, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30521 = stablehlo.multiply %30520, %30518 : tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30522 = stablehlo.multiply %30505, %30521 : tensor<1x1x2048x8192xbf16> loc(#loc581)
      %30523 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc584)
      %30524 = stablehlo.convert %30523 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc585)
      %30525 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30526 = stablehlo.broadcast_in_dim %30525, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc586)
      %30527 = stablehlo.subtract %30526, %30524 : tensor<1x1x2048x1xf32> loc(#loc586)
      %30528 = stablehlo.convert %30527 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc587)
      %30529 = stablehlo.broadcast_in_dim %30528, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc588)
      %30530 = stablehlo.multiply %30522, %30529 : tensor<1x1x2048x8192xbf16> loc(#loc588)
      %30531 = func.call @_where_143(%30042, %iterArg_100, %iterArg_101) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %30532 = func.call @_einsum_144(%30530, %30531) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc589)
      %30533 = arith.select %ptrue, %30532, %30532 : tensor<1x1x2048x2048xbf16>
      %30534 = func.call @_where_145(%30042, %iterArg_102, %iterArg_103) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30535 = stablehlo.broadcast_in_dim %30534, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc591)
      %30536 = stablehlo.broadcast_in_dim %30535, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc592)
      %30537 = stablehlo.add %30533, %30536 : tensor<1x1x2048x2048xbf16> loc(#loc592)
      %30538 = stablehlo.convert %30523 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc585)
      %30539 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30540 = stablehlo.broadcast_in_dim %30539, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc586)
      %30541 = stablehlo.subtract %30540, %30538 : tensor<1x1x2048x1xf32> loc(#loc586)
      %30542 = stablehlo.convert %30541 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc593)
      %30543 = stablehlo.broadcast_in_dim %30542, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc594)
      %30544 = stablehlo.multiply %30537, %30543 : tensor<1x1x2048x2048xbf16> loc(#loc594)
      %30545 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30546 = stablehlo.broadcast_in_dim %30545, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc595)
      %30547 = stablehlo.multiply %30544, %30546 : tensor<1x1x2048x2048xbf16> loc(#loc595)
      %30548 = stablehlo.add %30460, %30547 : tensor<1x1x2048x2048xbf16> loc(#loc596)
      %30549 = stablehlo.convert %30548 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc597)
      %30550 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc598)
      %30551 = stablehlo.reduce(%30549 init: %30550) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc598)
        stablehlo.return %33956 : tensor<f32> loc(#loc598)
      } loc(#loc598)
      %30552 = stablehlo.broadcast_in_dim %30551, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc599)
      %30553 = stablehlo.convert %30552 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc600)
      %30554 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30555 = stablehlo.broadcast_in_dim %30554, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc601)
      %30556 = stablehlo.divide %30553, %30555 : tensor<1x1x2048x1xbf16> loc(#loc601)
      %30557 = stablehlo.broadcast_in_dim %30556, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc602)
      %30558 = stablehlo.subtract %30548, %30557 : tensor<1x1x2048x2048xbf16> loc(#loc602)
      %30559 = stablehlo.broadcast_in_dim %30556, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc602)
      %30560 = stablehlo.subtract %30548, %30559 : tensor<1x1x2048x2048xbf16> loc(#loc602)
      %30561 = stablehlo.multiply %30560, %30560 : tensor<1x1x2048x2048xbf16> loc(#loc603)
      %30562 = stablehlo.convert %30561 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc597)
      %30563 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc598)
      %30564 = stablehlo.reduce(%30562 init: %30563) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc598)
        stablehlo.return %33956 : tensor<f32> loc(#loc598)
      } loc(#loc598)
      %30565 = stablehlo.broadcast_in_dim %30564, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc599)
      %30566 = stablehlo.convert %30565 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc600)
      %30567 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30568 = stablehlo.broadcast_in_dim %30567, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc601)
      %30569 = stablehlo.divide %30566, %30568 : tensor<1x1x2048x1xbf16> loc(#loc601)
      %30570 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30571 = stablehlo.broadcast_in_dim %30570, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc604)
      %30572 = stablehlo.add %30569, %30571 : tensor<1x1x2048x1xbf16> loc(#loc604)
      %30573 = stablehlo.rsqrt %30572 : tensor<1x1x2048x1xbf16> loc(#loc605)
      %30574 = stablehlo.broadcast_in_dim %30573, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc606)
      %30575 = stablehlo.multiply %30558, %30574 : tensor<1x1x2048x2048xbf16> loc(#loc606)
      %30576 = func.call @_where_146(%30042, %iterArg_104, %iterArg_105) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30577 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30578 = stablehlo.broadcast_in_dim %30577, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc607)
      %30579 = stablehlo.add %30578, %30576 : tensor<1x2048xbf16> loc(#loc607)
      %30580 = stablehlo.broadcast_in_dim %30579, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc608)
      %30581 = stablehlo.broadcast_in_dim %30580, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc609)
      %30582 = stablehlo.multiply %30575, %30581 : tensor<1x1x2048x2048xbf16> loc(#loc609)
      %30583 = func.call @_where_147(%30042, %iterArg_106, %iterArg_107) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30584 = stablehlo.broadcast_in_dim %30583, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc610)
      %30585 = stablehlo.broadcast_in_dim %30584, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc611)
      %30586 = stablehlo.add %30582, %30585 : tensor<1x1x2048x2048xbf16> loc(#loc611)
      %30587 = func.call @_where_148(%30042, %iterArg_108, %iterArg_109) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %30588 = func.call @_einsum_149(%30586, %30587) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc612)
      %30589 = "stablehlo.slice"(%30588) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc613)
      %30590 = stablehlo.reshape %30589 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc614)
      %30591 = arith.select %ptrue, %30590, %30590 : tensor<1x1x2048x32x64xbf16>
      %30592 = func.call @_where_150(%30042, %iterArg_110, %iterArg_111) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %30593 = func.call @softplus_151(%30592, %iterArg_112) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc616)
      %30594 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %30595 = stablehlo.broadcast_in_dim %30594, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc617)
      %30596 = stablehlo.multiply %30595, %30593 : tensor<1x64xbf16> loc(#loc617)
      %30597 = stablehlo.broadcast_in_dim %30596, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc618)
      %30598 = stablehlo.broadcast_in_dim %30597, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc619)
      %30599 = stablehlo.multiply %30591, %30598 : tensor<1x1x2048x32x64xbf16> loc(#loc619)
      %30600 = "stablehlo.slice"(%30588) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc620)
      %30601 = stablehlo.reshape %30600 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc614)
      %30602 = arith.select %ptrue, %30601, %30601 : tensor<1x1x2048x32x64xbf16>
      %30603 = func.call @_einsum_153(%30599, %30602) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc621)
      %30604 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30605 = stablehlo.broadcast_in_dim %30604, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc622)
      %30606 = stablehlo.divide %30603, %30605 : tensor<1x1x32x2048x2048xbf16> loc(#loc622)
      %30607 = stablehlo.tanh %30606 : tensor<1x1x32x2048x2048xbf16> loc(#loc623)
      %30608 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30609 = stablehlo.broadcast_in_dim %30608, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc624)
      %30610 = stablehlo.multiply %30609, %30607 : tensor<1x1x32x2048x2048xbf16> loc(#loc624)
      %30611 = stablehlo.convert %30610 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc625)
      %30612 = func.call @_where_154(%29988, %30611, %iterArg_113) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc626)
      %30613 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc627)
      %30614 = stablehlo.reduce(%30612 init: %30613) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc627)
        stablehlo.return %33956 : tensor<f32> loc(#loc627)
      } loc(#loc627)
      %30615 = stablehlo.broadcast_in_dim %30614, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc628)
      %30616 = stablehlo.broadcast_in_dim %30615, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc629)
      %30617 = stablehlo.subtract %30612, %30616 : tensor<1x1x32x2048x2048xf32> loc(#loc629)
      %30618 = stablehlo.exponential %30617 : tensor<1x1x32x2048x2048xf32> loc(#loc630)
      %30619 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc631)
      %30620 = stablehlo.reduce(%30618 init: %30619) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc631)
        stablehlo.return %33956 : tensor<f32> loc(#loc631)
      } loc(#loc631)
      %30621 = stablehlo.broadcast_in_dim %30620, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc628)
      %30622 = stablehlo.broadcast_in_dim %30621, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc632)
      %30623 = stablehlo.divide %30618, %30622 : tensor<1x1x32x2048x2048xf32> loc(#loc632)
      %30624 = stablehlo.convert %30623 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc633)
      %30625 = "stablehlo.slice"(%30588) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc634)
      %30626 = stablehlo.reshape %30625 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc614)
      %30627 = arith.select %ptrue, %30626, %30626 : tensor<1x1x2048x32x64xbf16>
      %30628 = func.call @_einsum_155(%30624, %30627) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc635)
      %30629 = arith.select %ptrue, %30628, %30628 : tensor<1x1x2048x32x64xbf16>
      %30630 = func.call @_where_156(%30042, %iterArg_114, %iterArg_115) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %30631 = func.call @_einsum_157(%30629, %30630) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc636)
      %30632 = arith.select %ptrue, %30631, %30631 : tensor<1x1x2048x2048xbf16>
      %30633 = stablehlo.add %30632, %30548 : tensor<1x1x2048x2048xbf16> loc(#loc638)
      %30634 = stablehlo.convert %30633 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc639)
      %30635 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc640)
      %30636 = stablehlo.reduce(%30634 init: %30635) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc640)
        stablehlo.return %33956 : tensor<f32> loc(#loc640)
      } loc(#loc640)
      %30637 = stablehlo.broadcast_in_dim %30636, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc641)
      %30638 = stablehlo.convert %30637 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc642)
      %30639 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30640 = stablehlo.broadcast_in_dim %30639, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc643)
      %30641 = stablehlo.divide %30638, %30640 : tensor<1x1x2048x1xbf16> loc(#loc643)
      %30642 = stablehlo.broadcast_in_dim %30641, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc644)
      %30643 = stablehlo.subtract %30633, %30642 : tensor<1x1x2048x2048xbf16> loc(#loc644)
      %30644 = stablehlo.broadcast_in_dim %30641, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc644)
      %30645 = stablehlo.subtract %30633, %30644 : tensor<1x1x2048x2048xbf16> loc(#loc644)
      %30646 = stablehlo.multiply %30645, %30645 : tensor<1x1x2048x2048xbf16> loc(#loc645)
      %30647 = stablehlo.convert %30646 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc639)
      %30648 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc640)
      %30649 = stablehlo.reduce(%30647 init: %30648) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc640)
        stablehlo.return %33956 : tensor<f32> loc(#loc640)
      } loc(#loc640)
      %30650 = stablehlo.broadcast_in_dim %30649, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc641)
      %30651 = stablehlo.convert %30650 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc642)
      %30652 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30653 = stablehlo.broadcast_in_dim %30652, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc643)
      %30654 = stablehlo.divide %30651, %30653 : tensor<1x1x2048x1xbf16> loc(#loc643)
      %30655 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30656 = stablehlo.broadcast_in_dim %30655, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc646)
      %30657 = stablehlo.add %30654, %30656 : tensor<1x1x2048x1xbf16> loc(#loc646)
      %30658 = stablehlo.rsqrt %30657 : tensor<1x1x2048x1xbf16> loc(#loc647)
      %30659 = stablehlo.broadcast_in_dim %30658, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc648)
      %30660 = stablehlo.multiply %30643, %30659 : tensor<1x1x2048x2048xbf16> loc(#loc648)
      %30661 = func.call @_where_158(%30042, %iterArg_116, %iterArg_117) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30662 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30663 = stablehlo.broadcast_in_dim %30662, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc649)
      %30664 = stablehlo.add %30663, %30661 : tensor<1x2048xbf16> loc(#loc649)
      %30665 = stablehlo.broadcast_in_dim %30664, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc650)
      %30666 = stablehlo.broadcast_in_dim %30665, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc651)
      %30667 = stablehlo.multiply %30660, %30666 : tensor<1x1x2048x2048xbf16> loc(#loc651)
      %30668 = func.call @_where_159(%30042, %iterArg_118, %iterArg_119) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30669 = stablehlo.broadcast_in_dim %30668, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc652)
      %30670 = stablehlo.broadcast_in_dim %30669, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc653)
      %30671 = stablehlo.add %30667, %30670 : tensor<1x1x2048x2048xbf16> loc(#loc653)
      %30672 = func.call @_where_160(%30042, %iterArg_120, %iterArg_121) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %30673 = func.call @_einsum_161(%30671, %30672) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc654)
      %30674 = arith.select %ptrue, %30673, %30673 : tensor<1x1x2048x8192xbf16>
      %30675 = func.call @_where_162(%30042, %iterArg_122, %iterArg_123) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %30676 = stablehlo.broadcast_in_dim %30675, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc656)
      %30677 = stablehlo.broadcast_in_dim %30676, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc657)
      %30678 = stablehlo.add %30674, %30677 : tensor<1x1x2048x8192xbf16> loc(#loc657)
      %30679 = stablehlo.multiply %30678, %30678 : tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30680 = stablehlo.multiply %30678, %30679 : tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30681 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %30682 = stablehlo.broadcast_in_dim %30681, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30683 = stablehlo.multiply %30682, %30680 : tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30684 = stablehlo.add %30678, %30683 : tensor<1x1x2048x8192xbf16> loc(#loc659)
      %30685 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %30686 = stablehlo.broadcast_in_dim %30685, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30687 = stablehlo.multiply %30686, %30684 : tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30688 = stablehlo.tanh %30687 : tensor<1x1x2048x8192xbf16> loc(#loc660)
      %30689 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30690 = stablehlo.broadcast_in_dim %30689, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc659)
      %30691 = stablehlo.add %30690, %30688 : tensor<1x1x2048x8192xbf16> loc(#loc659)
      %30692 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %30693 = stablehlo.broadcast_in_dim %30692, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30694 = stablehlo.multiply %30693, %30691 : tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30695 = stablehlo.multiply %30678, %30694 : tensor<1x1x2048x8192xbf16> loc(#loc658)
      %30696 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc661)
      %30697 = stablehlo.convert %30696 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc662)
      %30698 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30699 = stablehlo.broadcast_in_dim %30698, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc663)
      %30700 = stablehlo.subtract %30699, %30697 : tensor<1x1x2048x1xf32> loc(#loc663)
      %30701 = stablehlo.convert %30700 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc664)
      %30702 = stablehlo.broadcast_in_dim %30701, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc665)
      %30703 = stablehlo.multiply %30695, %30702 : tensor<1x1x2048x8192xbf16> loc(#loc665)
      %30704 = func.call @_where_163(%30042, %iterArg_124, %iterArg_125) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %30705 = func.call @_einsum_164(%30703, %30704) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc666)
      %30706 = arith.select %ptrue, %30705, %30705 : tensor<1x1x2048x2048xbf16>
      %30707 = func.call @_where_165(%30042, %iterArg_126, %iterArg_127) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30708 = stablehlo.broadcast_in_dim %30707, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc668)
      %30709 = stablehlo.broadcast_in_dim %30708, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc669)
      %30710 = stablehlo.add %30706, %30709 : tensor<1x1x2048x2048xbf16> loc(#loc669)
      %30711 = stablehlo.convert %30696 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc662)
      %30712 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30713 = stablehlo.broadcast_in_dim %30712, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc663)
      %30714 = stablehlo.subtract %30713, %30711 : tensor<1x1x2048x1xf32> loc(#loc663)
      %30715 = stablehlo.convert %30714 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc670)
      %30716 = stablehlo.broadcast_in_dim %30715, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc671)
      %30717 = stablehlo.multiply %30710, %30716 : tensor<1x1x2048x2048xbf16> loc(#loc671)
      %30718 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30719 = stablehlo.broadcast_in_dim %30718, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc672)
      %30720 = stablehlo.multiply %30717, %30719 : tensor<1x1x2048x2048xbf16> loc(#loc672)
      %30721 = stablehlo.add %30633, %30720 : tensor<1x1x2048x2048xbf16> loc(#loc673)
      %30722 = stablehlo.convert %30721 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc674)
      %30723 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc675)
      %30724 = stablehlo.reduce(%30722 init: %30723) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc675)
        stablehlo.return %33956 : tensor<f32> loc(#loc675)
      } loc(#loc675)
      %30725 = stablehlo.broadcast_in_dim %30724, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc676)
      %30726 = stablehlo.convert %30725 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc677)
      %30727 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30728 = stablehlo.broadcast_in_dim %30727, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc678)
      %30729 = stablehlo.divide %30726, %30728 : tensor<1x1x2048x1xbf16> loc(#loc678)
      %30730 = stablehlo.broadcast_in_dim %30729, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc679)
      %30731 = stablehlo.subtract %30721, %30730 : tensor<1x1x2048x2048xbf16> loc(#loc679)
      %30732 = stablehlo.broadcast_in_dim %30729, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc679)
      %30733 = stablehlo.subtract %30721, %30732 : tensor<1x1x2048x2048xbf16> loc(#loc679)
      %30734 = stablehlo.multiply %30733, %30733 : tensor<1x1x2048x2048xbf16> loc(#loc680)
      %30735 = stablehlo.convert %30734 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc674)
      %30736 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc675)
      %30737 = stablehlo.reduce(%30735 init: %30736) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc675)
        stablehlo.return %33956 : tensor<f32> loc(#loc675)
      } loc(#loc675)
      %30738 = stablehlo.broadcast_in_dim %30737, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc676)
      %30739 = stablehlo.convert %30738 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc677)
      %30740 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30741 = stablehlo.broadcast_in_dim %30740, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc678)
      %30742 = stablehlo.divide %30739, %30741 : tensor<1x1x2048x1xbf16> loc(#loc678)
      %30743 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30744 = stablehlo.broadcast_in_dim %30743, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc681)
      %30745 = stablehlo.add %30742, %30744 : tensor<1x1x2048x1xbf16> loc(#loc681)
      %30746 = stablehlo.rsqrt %30745 : tensor<1x1x2048x1xbf16> loc(#loc682)
      %30747 = stablehlo.broadcast_in_dim %30746, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc683)
      %30748 = stablehlo.multiply %30731, %30747 : tensor<1x1x2048x2048xbf16> loc(#loc683)
      %30749 = func.call @_where_166(%30042, %iterArg_128, %iterArg_129) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30750 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30751 = stablehlo.broadcast_in_dim %30750, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc684)
      %30752 = stablehlo.add %30751, %30749 : tensor<1x2048xbf16> loc(#loc684)
      %30753 = stablehlo.broadcast_in_dim %30752, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc685)
      %30754 = stablehlo.broadcast_in_dim %30753, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc686)
      %30755 = stablehlo.multiply %30748, %30754 : tensor<1x1x2048x2048xbf16> loc(#loc686)
      %30756 = func.call @_where_167(%30042, %iterArg_130, %iterArg_131) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30757 = stablehlo.broadcast_in_dim %30756, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc687)
      %30758 = stablehlo.broadcast_in_dim %30757, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc688)
      %30759 = stablehlo.add %30755, %30758 : tensor<1x1x2048x2048xbf16> loc(#loc688)
      %30760 = func.call @_where_168(%30042, %iterArg_132, %iterArg_133) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %30761 = func.call @_einsum_169(%30759, %30760) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc689)
      %30762 = "stablehlo.slice"(%30761) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc690)
      %30763 = stablehlo.reshape %30762 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc691)
      %30764 = arith.select %ptrue, %30763, %30763 : tensor<1x1x2048x32x64xbf16>
      %30765 = func.call @_where_170(%30042, %iterArg_134, %iterArg_135) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %30766 = func.call @softplus_171(%30765, %iterArg_136) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc693)
      %30767 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %30768 = stablehlo.broadcast_in_dim %30767, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc694)
      %30769 = stablehlo.multiply %30768, %30766 : tensor<1x64xbf16> loc(#loc694)
      %30770 = stablehlo.broadcast_in_dim %30769, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc695)
      %30771 = stablehlo.broadcast_in_dim %30770, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc696)
      %30772 = stablehlo.multiply %30764, %30771 : tensor<1x1x2048x32x64xbf16> loc(#loc696)
      %30773 = "stablehlo.slice"(%30761) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc697)
      %30774 = stablehlo.reshape %30773 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc691)
      %30775 = arith.select %ptrue, %30774, %30774 : tensor<1x1x2048x32x64xbf16>
      %30776 = func.call @_einsum_173(%30772, %30775) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc698)
      %30777 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30778 = stablehlo.broadcast_in_dim %30777, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc699)
      %30779 = stablehlo.divide %30776, %30778 : tensor<1x1x32x2048x2048xbf16> loc(#loc699)
      %30780 = stablehlo.tanh %30779 : tensor<1x1x32x2048x2048xbf16> loc(#loc700)
      %30781 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30782 = stablehlo.broadcast_in_dim %30781, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc701)
      %30783 = stablehlo.multiply %30782, %30780 : tensor<1x1x32x2048x2048xbf16> loc(#loc701)
      %30784 = stablehlo.convert %30783 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc702)
      %30785 = func.call @_where_174(%29985, %30784, %iterArg_137) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc703)
      %30786 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc704)
      %30787 = stablehlo.reduce(%30785 init: %30786) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc704)
        stablehlo.return %33956 : tensor<f32> loc(#loc704)
      } loc(#loc704)
      %30788 = stablehlo.broadcast_in_dim %30787, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc705)
      %30789 = stablehlo.broadcast_in_dim %30788, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc706)
      %30790 = stablehlo.subtract %30785, %30789 : tensor<1x1x32x2048x2048xf32> loc(#loc706)
      %30791 = stablehlo.exponential %30790 : tensor<1x1x32x2048x2048xf32> loc(#loc707)
      %30792 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc708)
      %30793 = stablehlo.reduce(%30791 init: %30792) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc708)
        stablehlo.return %33956 : tensor<f32> loc(#loc708)
      } loc(#loc708)
      %30794 = stablehlo.broadcast_in_dim %30793, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc705)
      %30795 = stablehlo.broadcast_in_dim %30794, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc709)
      %30796 = stablehlo.divide %30791, %30795 : tensor<1x1x32x2048x2048xf32> loc(#loc709)
      %30797 = stablehlo.convert %30796 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc710)
      %30798 = "stablehlo.slice"(%30761) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc711)
      %30799 = stablehlo.reshape %30798 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc691)
      %30800 = arith.select %ptrue, %30799, %30799 : tensor<1x1x2048x32x64xbf16>
      %30801 = func.call @_einsum_175(%30797, %30800) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc712)
      %30802 = arith.select %ptrue, %30801, %30801 : tensor<1x1x2048x32x64xbf16>
      %30803 = func.call @_where_176(%30042, %iterArg_138, %iterArg_139) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %30804 = func.call @_einsum_177(%30802, %30803) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc713)
      %30805 = arith.select %ptrue, %30804, %30804 : tensor<1x1x2048x2048xbf16>
      %30806 = stablehlo.add %30805, %30721 : tensor<1x1x2048x2048xbf16> loc(#loc715)
      %30807 = stablehlo.convert %30806 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc716)
      %30808 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc717)
      %30809 = stablehlo.reduce(%30807 init: %30808) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc717)
        stablehlo.return %33956 : tensor<f32> loc(#loc717)
      } loc(#loc717)
      %30810 = stablehlo.broadcast_in_dim %30809, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc718)
      %30811 = stablehlo.convert %30810 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc719)
      %30812 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30813 = stablehlo.broadcast_in_dim %30812, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc720)
      %30814 = stablehlo.divide %30811, %30813 : tensor<1x1x2048x1xbf16> loc(#loc720)
      %30815 = stablehlo.broadcast_in_dim %30814, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc721)
      %30816 = stablehlo.subtract %30806, %30815 : tensor<1x1x2048x2048xbf16> loc(#loc721)
      %30817 = stablehlo.broadcast_in_dim %30814, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc721)
      %30818 = stablehlo.subtract %30806, %30817 : tensor<1x1x2048x2048xbf16> loc(#loc721)
      %30819 = stablehlo.multiply %30818, %30818 : tensor<1x1x2048x2048xbf16> loc(#loc722)
      %30820 = stablehlo.convert %30819 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc716)
      %30821 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc717)
      %30822 = stablehlo.reduce(%30820 init: %30821) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc717)
        stablehlo.return %33956 : tensor<f32> loc(#loc717)
      } loc(#loc717)
      %30823 = stablehlo.broadcast_in_dim %30822, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc718)
      %30824 = stablehlo.convert %30823 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc719)
      %30825 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30826 = stablehlo.broadcast_in_dim %30825, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc720)
      %30827 = stablehlo.divide %30824, %30826 : tensor<1x1x2048x1xbf16> loc(#loc720)
      %30828 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30829 = stablehlo.broadcast_in_dim %30828, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc723)
      %30830 = stablehlo.add %30827, %30829 : tensor<1x1x2048x1xbf16> loc(#loc723)
      %30831 = stablehlo.rsqrt %30830 : tensor<1x1x2048x1xbf16> loc(#loc724)
      %30832 = stablehlo.broadcast_in_dim %30831, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc725)
      %30833 = stablehlo.multiply %30816, %30832 : tensor<1x1x2048x2048xbf16> loc(#loc725)
      %30834 = func.call @_where_178(%30042, %iterArg_140, %iterArg_141) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30835 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30836 = stablehlo.broadcast_in_dim %30835, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc726)
      %30837 = stablehlo.add %30836, %30834 : tensor<1x2048xbf16> loc(#loc726)
      %30838 = stablehlo.broadcast_in_dim %30837, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc727)
      %30839 = stablehlo.broadcast_in_dim %30838, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc728)
      %30840 = stablehlo.multiply %30833, %30839 : tensor<1x1x2048x2048xbf16> loc(#loc728)
      %30841 = func.call @_where_179(%30042, %iterArg_142, %iterArg_143) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30842 = stablehlo.broadcast_in_dim %30841, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc729)
      %30843 = stablehlo.broadcast_in_dim %30842, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc730)
      %30844 = stablehlo.add %30840, %30843 : tensor<1x1x2048x2048xbf16> loc(#loc730)
      %30845 = func.call @_where_180(%30042, %iterArg_144, %iterArg_145) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %30846 = func.call @_einsum_181(%30844, %30845) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc731)
      %30847 = arith.select %ptrue, %30846, %30846 : tensor<1x1x2048x8192xbf16>
      %30848 = func.call @_where_182(%30042, %iterArg_146, %iterArg_147) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %30849 = stablehlo.broadcast_in_dim %30848, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc733)
      %30850 = stablehlo.broadcast_in_dim %30849, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc734)
      %30851 = stablehlo.add %30847, %30850 : tensor<1x1x2048x8192xbf16> loc(#loc734)
      %30852 = stablehlo.multiply %30851, %30851 : tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30853 = stablehlo.multiply %30851, %30852 : tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30854 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %30855 = stablehlo.broadcast_in_dim %30854, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30856 = stablehlo.multiply %30855, %30853 : tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30857 = stablehlo.add %30851, %30856 : tensor<1x1x2048x8192xbf16> loc(#loc736)
      %30858 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %30859 = stablehlo.broadcast_in_dim %30858, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30860 = stablehlo.multiply %30859, %30857 : tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30861 = stablehlo.tanh %30860 : tensor<1x1x2048x8192xbf16> loc(#loc737)
      %30862 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30863 = stablehlo.broadcast_in_dim %30862, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc736)
      %30864 = stablehlo.add %30863, %30861 : tensor<1x1x2048x8192xbf16> loc(#loc736)
      %30865 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %30866 = stablehlo.broadcast_in_dim %30865, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30867 = stablehlo.multiply %30866, %30864 : tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30868 = stablehlo.multiply %30851, %30867 : tensor<1x1x2048x8192xbf16> loc(#loc735)
      %30869 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc738)
      %30870 = stablehlo.convert %30869 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc739)
      %30871 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30872 = stablehlo.broadcast_in_dim %30871, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc740)
      %30873 = stablehlo.subtract %30872, %30870 : tensor<1x1x2048x1xf32> loc(#loc740)
      %30874 = stablehlo.convert %30873 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc741)
      %30875 = stablehlo.broadcast_in_dim %30874, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc742)
      %30876 = stablehlo.multiply %30868, %30875 : tensor<1x1x2048x8192xbf16> loc(#loc742)
      %30877 = func.call @_where_183(%30042, %iterArg_148, %iterArg_149) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %30878 = func.call @_einsum_184(%30876, %30877) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc743)
      %30879 = arith.select %ptrue, %30878, %30878 : tensor<1x1x2048x2048xbf16>
      %30880 = func.call @_where_185(%30042, %iterArg_150, %iterArg_151) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30881 = stablehlo.broadcast_in_dim %30880, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc745)
      %30882 = stablehlo.broadcast_in_dim %30881, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc746)
      %30883 = stablehlo.add %30879, %30882 : tensor<1x1x2048x2048xbf16> loc(#loc746)
      %30884 = stablehlo.convert %30869 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc739)
      %30885 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %30886 = stablehlo.broadcast_in_dim %30885, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc740)
      %30887 = stablehlo.subtract %30886, %30884 : tensor<1x1x2048x1xf32> loc(#loc740)
      %30888 = stablehlo.convert %30887 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc747)
      %30889 = stablehlo.broadcast_in_dim %30888, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc748)
      %30890 = stablehlo.multiply %30883, %30889 : tensor<1x1x2048x2048xbf16> loc(#loc748)
      %30891 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30892 = stablehlo.broadcast_in_dim %30891, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc749)
      %30893 = stablehlo.multiply %30890, %30892 : tensor<1x1x2048x2048xbf16> loc(#loc749)
      %30894 = stablehlo.add %30806, %30893 : tensor<1x1x2048x2048xbf16> loc(#loc750)
      %30895 = stablehlo.convert %30894 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc751)
      %30896 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc752)
      %30897 = stablehlo.reduce(%30895 init: %30896) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc752)
        stablehlo.return %33956 : tensor<f32> loc(#loc752)
      } loc(#loc752)
      %30898 = stablehlo.broadcast_in_dim %30897, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc753)
      %30899 = stablehlo.convert %30898 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc754)
      %30900 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30901 = stablehlo.broadcast_in_dim %30900, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc755)
      %30902 = stablehlo.divide %30899, %30901 : tensor<1x1x2048x1xbf16> loc(#loc755)
      %30903 = stablehlo.broadcast_in_dim %30902, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc756)
      %30904 = stablehlo.subtract %30894, %30903 : tensor<1x1x2048x2048xbf16> loc(#loc756)
      %30905 = stablehlo.broadcast_in_dim %30902, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc756)
      %30906 = stablehlo.subtract %30894, %30905 : tensor<1x1x2048x2048xbf16> loc(#loc756)
      %30907 = stablehlo.multiply %30906, %30906 : tensor<1x1x2048x2048xbf16> loc(#loc757)
      %30908 = stablehlo.convert %30907 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc751)
      %30909 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc752)
      %30910 = stablehlo.reduce(%30908 init: %30909) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc752)
        stablehlo.return %33956 : tensor<f32> loc(#loc752)
      } loc(#loc752)
      %30911 = stablehlo.broadcast_in_dim %30910, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc753)
      %30912 = stablehlo.convert %30911 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc754)
      %30913 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30914 = stablehlo.broadcast_in_dim %30913, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc755)
      %30915 = stablehlo.divide %30912, %30914 : tensor<1x1x2048x1xbf16> loc(#loc755)
      %30916 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %30917 = stablehlo.broadcast_in_dim %30916, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc758)
      %30918 = stablehlo.add %30915, %30917 : tensor<1x1x2048x1xbf16> loc(#loc758)
      %30919 = stablehlo.rsqrt %30918 : tensor<1x1x2048x1xbf16> loc(#loc759)
      %30920 = stablehlo.broadcast_in_dim %30919, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc760)
      %30921 = stablehlo.multiply %30904, %30920 : tensor<1x1x2048x2048xbf16> loc(#loc760)
      %30922 = func.call @_where_186(%30042, %iterArg_152, %iterArg_153) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30923 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %30924 = stablehlo.broadcast_in_dim %30923, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc761)
      %30925 = stablehlo.add %30924, %30922 : tensor<1x2048xbf16> loc(#loc761)
      %30926 = stablehlo.broadcast_in_dim %30925, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc762)
      %30927 = stablehlo.broadcast_in_dim %30926, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc763)
      %30928 = stablehlo.multiply %30921, %30927 : tensor<1x1x2048x2048xbf16> loc(#loc763)
      %30929 = func.call @_where_187(%30042, %iterArg_154, %iterArg_155) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %30930 = stablehlo.broadcast_in_dim %30929, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc764)
      %30931 = stablehlo.broadcast_in_dim %30930, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc765)
      %30932 = stablehlo.add %30928, %30931 : tensor<1x1x2048x2048xbf16> loc(#loc765)
      %30933 = func.call @_where_188(%30042, %iterArg_156, %iterArg_157) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %30934 = func.call @_einsum_189(%30932, %30933) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc766)
      %30935 = "stablehlo.slice"(%30934) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc767)
      %30936 = stablehlo.reshape %30935 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc768)
      %30937 = arith.select %ptrue, %30936, %30936 : tensor<1x1x2048x32x64xbf16>
      %30938 = func.call @_where_190(%30042, %iterArg_158, %iterArg_159) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %30939 = func.call @softplus_191(%30938, %iterArg_160) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc770)
      %30940 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %30941 = stablehlo.broadcast_in_dim %30940, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc771)
      %30942 = stablehlo.multiply %30941, %30939 : tensor<1x64xbf16> loc(#loc771)
      %30943 = stablehlo.broadcast_in_dim %30942, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc772)
      %30944 = stablehlo.broadcast_in_dim %30943, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc773)
      %30945 = stablehlo.multiply %30937, %30944 : tensor<1x1x2048x32x64xbf16> loc(#loc773)
      %30946 = "stablehlo.slice"(%30934) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc774)
      %30947 = stablehlo.reshape %30946 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc768)
      %30948 = arith.select %ptrue, %30947, %30947 : tensor<1x1x2048x32x64xbf16>
      %30949 = func.call @_einsum_193(%30945, %30948) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc775)
      %30950 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30951 = stablehlo.broadcast_in_dim %30950, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc776)
      %30952 = stablehlo.divide %30949, %30951 : tensor<1x1x32x2048x2048xbf16> loc(#loc776)
      %30953 = stablehlo.tanh %30952 : tensor<1x1x32x2048x2048xbf16> loc(#loc777)
      %30954 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %30955 = stablehlo.broadcast_in_dim %30954, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc778)
      %30956 = stablehlo.multiply %30955, %30953 : tensor<1x1x32x2048x2048xbf16> loc(#loc778)
      %30957 = stablehlo.convert %30956 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc779)
      %30958 = func.call @_where_194(%29982, %30957, %iterArg_161) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc780)
      %30959 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc781)
      %30960 = stablehlo.reduce(%30958 init: %30959) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc781)
        stablehlo.return %33956 : tensor<f32> loc(#loc781)
      } loc(#loc781)
      %30961 = stablehlo.broadcast_in_dim %30960, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc782)
      %30962 = stablehlo.broadcast_in_dim %30961, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc783)
      %30963 = stablehlo.subtract %30958, %30962 : tensor<1x1x32x2048x2048xf32> loc(#loc783)
      %30964 = stablehlo.exponential %30963 : tensor<1x1x32x2048x2048xf32> loc(#loc784)
      %30965 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc785)
      %30966 = stablehlo.reduce(%30964 init: %30965) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc785)
        stablehlo.return %33956 : tensor<f32> loc(#loc785)
      } loc(#loc785)
      %30967 = stablehlo.broadcast_in_dim %30966, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc782)
      %30968 = stablehlo.broadcast_in_dim %30967, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc786)
      %30969 = stablehlo.divide %30964, %30968 : tensor<1x1x32x2048x2048xf32> loc(#loc786)
      %30970 = stablehlo.convert %30969 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc787)
      %30971 = "stablehlo.slice"(%30934) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc788)
      %30972 = stablehlo.reshape %30971 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc768)
      %30973 = arith.select %ptrue, %30972, %30972 : tensor<1x1x2048x32x64xbf16>
      %30974 = func.call @_einsum_195(%30970, %30973) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc789)
      %30975 = arith.select %ptrue, %30974, %30974 : tensor<1x1x2048x32x64xbf16>
      %30976 = func.call @_where_196(%30042, %iterArg_162, %iterArg_163) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %30977 = func.call @_einsum_197(%30975, %30976) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc790)
      %30978 = arith.select %ptrue, %30977, %30977 : tensor<1x1x2048x2048xbf16>
      %30979 = stablehlo.add %30978, %30894 : tensor<1x1x2048x2048xbf16> loc(#loc792)
      %30980 = stablehlo.convert %30979 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc793)
      %30981 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc794)
      %30982 = stablehlo.reduce(%30980 init: %30981) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc794)
        stablehlo.return %33956 : tensor<f32> loc(#loc794)
      } loc(#loc794)
      %30983 = stablehlo.broadcast_in_dim %30982, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc795)
      %30984 = stablehlo.convert %30983 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc796)
      %30985 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30986 = stablehlo.broadcast_in_dim %30985, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc797)
      %30987 = stablehlo.divide %30984, %30986 : tensor<1x1x2048x1xbf16> loc(#loc797)
      %30988 = stablehlo.broadcast_in_dim %30987, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc798)
      %30989 = stablehlo.subtract %30979, %30988 : tensor<1x1x2048x2048xbf16> loc(#loc798)
      %30990 = stablehlo.broadcast_in_dim %30987, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc798)
      %30991 = stablehlo.subtract %30979, %30990 : tensor<1x1x2048x2048xbf16> loc(#loc798)
      %30992 = stablehlo.multiply %30991, %30991 : tensor<1x1x2048x2048xbf16> loc(#loc799)
      %30993 = stablehlo.convert %30992 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc793)
      %30994 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc794)
      %30995 = stablehlo.reduce(%30993 init: %30994) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc794)
        stablehlo.return %33956 : tensor<f32> loc(#loc794)
      } loc(#loc794)
      %30996 = stablehlo.broadcast_in_dim %30995, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc795)
      %30997 = stablehlo.convert %30996 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc796)
      %30998 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %30999 = stablehlo.broadcast_in_dim %30998, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc797)
      %31000 = stablehlo.divide %30997, %30999 : tensor<1x1x2048x1xbf16> loc(#loc797)
      %31001 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31002 = stablehlo.broadcast_in_dim %31001, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc800)
      %31003 = stablehlo.add %31000, %31002 : tensor<1x1x2048x1xbf16> loc(#loc800)
      %31004 = stablehlo.rsqrt %31003 : tensor<1x1x2048x1xbf16> loc(#loc801)
      %31005 = stablehlo.broadcast_in_dim %31004, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc802)
      %31006 = stablehlo.multiply %30989, %31005 : tensor<1x1x2048x2048xbf16> loc(#loc802)
      %31007 = func.call @_where_198(%30042, %iterArg_164, %iterArg_165) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31008 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31009 = stablehlo.broadcast_in_dim %31008, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc803)
      %31010 = stablehlo.add %31009, %31007 : tensor<1x2048xbf16> loc(#loc803)
      %31011 = stablehlo.broadcast_in_dim %31010, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc804)
      %31012 = stablehlo.broadcast_in_dim %31011, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc805)
      %31013 = stablehlo.multiply %31006, %31012 : tensor<1x1x2048x2048xbf16> loc(#loc805)
      %31014 = func.call @_where_199(%30042, %iterArg_166, %iterArg_167) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31015 = stablehlo.broadcast_in_dim %31014, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc806)
      %31016 = stablehlo.broadcast_in_dim %31015, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc807)
      %31017 = stablehlo.add %31013, %31016 : tensor<1x1x2048x2048xbf16> loc(#loc807)
      %31018 = func.call @_where_200(%30042, %iterArg_168, %iterArg_169) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %31019 = func.call @_einsum_201(%31017, %31018) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc808)
      %31020 = arith.select %ptrue, %31019, %31019 : tensor<1x1x2048x8192xbf16>
      %31021 = func.call @_where_202(%30042, %iterArg_170, %iterArg_171) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %31022 = stablehlo.broadcast_in_dim %31021, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc810)
      %31023 = stablehlo.broadcast_in_dim %31022, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc811)
      %31024 = stablehlo.add %31020, %31023 : tensor<1x1x2048x8192xbf16> loc(#loc811)
      %31025 = stablehlo.multiply %31024, %31024 : tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31026 = stablehlo.multiply %31024, %31025 : tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31027 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %31028 = stablehlo.broadcast_in_dim %31027, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31029 = stablehlo.multiply %31028, %31026 : tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31030 = stablehlo.add %31024, %31029 : tensor<1x1x2048x8192xbf16> loc(#loc813)
      %31031 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %31032 = stablehlo.broadcast_in_dim %31031, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31033 = stablehlo.multiply %31032, %31030 : tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31034 = stablehlo.tanh %31033 : tensor<1x1x2048x8192xbf16> loc(#loc814)
      %31035 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31036 = stablehlo.broadcast_in_dim %31035, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc813)
      %31037 = stablehlo.add %31036, %31034 : tensor<1x1x2048x8192xbf16> loc(#loc813)
      %31038 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %31039 = stablehlo.broadcast_in_dim %31038, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31040 = stablehlo.multiply %31039, %31037 : tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31041 = stablehlo.multiply %31024, %31040 : tensor<1x1x2048x8192xbf16> loc(#loc812)
      %31042 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc815)
      %31043 = stablehlo.convert %31042 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc816)
      %31044 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31045 = stablehlo.broadcast_in_dim %31044, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc817)
      %31046 = stablehlo.subtract %31045, %31043 : tensor<1x1x2048x1xf32> loc(#loc817)
      %31047 = stablehlo.convert %31046 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc818)
      %31048 = stablehlo.broadcast_in_dim %31047, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc819)
      %31049 = stablehlo.multiply %31041, %31048 : tensor<1x1x2048x8192xbf16> loc(#loc819)
      %31050 = func.call @_where_203(%30042, %iterArg_172, %iterArg_173) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %31051 = func.call @_einsum_204(%31049, %31050) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc820)
      %31052 = arith.select %ptrue, %31051, %31051 : tensor<1x1x2048x2048xbf16>
      %31053 = func.call @_where_205(%30042, %iterArg_174, %iterArg_175) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31054 = stablehlo.broadcast_in_dim %31053, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc822)
      %31055 = stablehlo.broadcast_in_dim %31054, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc823)
      %31056 = stablehlo.add %31052, %31055 : tensor<1x1x2048x2048xbf16> loc(#loc823)
      %31057 = stablehlo.convert %31042 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc816)
      %31058 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31059 = stablehlo.broadcast_in_dim %31058, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc817)
      %31060 = stablehlo.subtract %31059, %31057 : tensor<1x1x2048x1xf32> loc(#loc817)
      %31061 = stablehlo.convert %31060 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc824)
      %31062 = stablehlo.broadcast_in_dim %31061, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc825)
      %31063 = stablehlo.multiply %31056, %31062 : tensor<1x1x2048x2048xbf16> loc(#loc825)
      %31064 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31065 = stablehlo.broadcast_in_dim %31064, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc826)
      %31066 = stablehlo.multiply %31063, %31065 : tensor<1x1x2048x2048xbf16> loc(#loc826)
      %31067 = stablehlo.add %30979, %31066 : tensor<1x1x2048x2048xbf16> loc(#loc827)
      %31068 = stablehlo.convert %31067 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc828)
      %31069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc829)
      %31070 = stablehlo.reduce(%31068 init: %31069) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc829)
        stablehlo.return %33956 : tensor<f32> loc(#loc829)
      } loc(#loc829)
      %31071 = stablehlo.broadcast_in_dim %31070, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc830)
      %31072 = stablehlo.convert %31071 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc831)
      %31073 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31074 = stablehlo.broadcast_in_dim %31073, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc832)
      %31075 = stablehlo.divide %31072, %31074 : tensor<1x1x2048x1xbf16> loc(#loc832)
      %31076 = stablehlo.broadcast_in_dim %31075, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc833)
      %31077 = stablehlo.subtract %31067, %31076 : tensor<1x1x2048x2048xbf16> loc(#loc833)
      %31078 = stablehlo.broadcast_in_dim %31075, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc833)
      %31079 = stablehlo.subtract %31067, %31078 : tensor<1x1x2048x2048xbf16> loc(#loc833)
      %31080 = stablehlo.multiply %31079, %31079 : tensor<1x1x2048x2048xbf16> loc(#loc834)
      %31081 = stablehlo.convert %31080 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc828)
      %31082 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc829)
      %31083 = stablehlo.reduce(%31081 init: %31082) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc829)
        stablehlo.return %33956 : tensor<f32> loc(#loc829)
      } loc(#loc829)
      %31084 = stablehlo.broadcast_in_dim %31083, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc830)
      %31085 = stablehlo.convert %31084 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc831)
      %31086 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31087 = stablehlo.broadcast_in_dim %31086, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc832)
      %31088 = stablehlo.divide %31085, %31087 : tensor<1x1x2048x1xbf16> loc(#loc832)
      %31089 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31090 = stablehlo.broadcast_in_dim %31089, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc835)
      %31091 = stablehlo.add %31088, %31090 : tensor<1x1x2048x1xbf16> loc(#loc835)
      %31092 = stablehlo.rsqrt %31091 : tensor<1x1x2048x1xbf16> loc(#loc836)
      %31093 = stablehlo.broadcast_in_dim %31092, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc837)
      %31094 = stablehlo.multiply %31077, %31093 : tensor<1x1x2048x2048xbf16> loc(#loc837)
      %31095 = func.call @_where_206(%30042, %iterArg_176, %iterArg_177) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31096 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31097 = stablehlo.broadcast_in_dim %31096, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc838)
      %31098 = stablehlo.add %31097, %31095 : tensor<1x2048xbf16> loc(#loc838)
      %31099 = stablehlo.broadcast_in_dim %31098, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc839)
      %31100 = stablehlo.broadcast_in_dim %31099, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc840)
      %31101 = stablehlo.multiply %31094, %31100 : tensor<1x1x2048x2048xbf16> loc(#loc840)
      %31102 = func.call @_where_207(%30042, %iterArg_178, %iterArg_179) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31103 = stablehlo.broadcast_in_dim %31102, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc841)
      %31104 = stablehlo.broadcast_in_dim %31103, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc842)
      %31105 = stablehlo.add %31101, %31104 : tensor<1x1x2048x2048xbf16> loc(#loc842)
      %31106 = func.call @_where_208(%30042, %iterArg_180, %iterArg_181) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %31107 = func.call @_einsum_209(%31105, %31106) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc843)
      %31108 = "stablehlo.slice"(%31107) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc844)
      %31109 = stablehlo.reshape %31108 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc845)
      %31110 = arith.select %ptrue, %31109, %31109 : tensor<1x1x2048x32x64xbf16>
      %31111 = func.call @_where_210(%30042, %iterArg_182, %iterArg_183) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %31112 = func.call @softplus_211(%31111, %iterArg_184) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc847)
      %31113 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %31114 = stablehlo.broadcast_in_dim %31113, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc848)
      %31115 = stablehlo.multiply %31114, %31112 : tensor<1x64xbf16> loc(#loc848)
      %31116 = stablehlo.broadcast_in_dim %31115, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc849)
      %31117 = stablehlo.broadcast_in_dim %31116, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc850)
      %31118 = stablehlo.multiply %31110, %31117 : tensor<1x1x2048x32x64xbf16> loc(#loc850)
      %31119 = "stablehlo.slice"(%31107) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc851)
      %31120 = stablehlo.reshape %31119 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc845)
      %31121 = arith.select %ptrue, %31120, %31120 : tensor<1x1x2048x32x64xbf16>
      %31122 = func.call @_einsum_213(%31118, %31121) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc852)
      %31123 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31124 = stablehlo.broadcast_in_dim %31123, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc853)
      %31125 = stablehlo.divide %31122, %31124 : tensor<1x1x32x2048x2048xbf16> loc(#loc853)
      %31126 = stablehlo.tanh %31125 : tensor<1x1x32x2048x2048xbf16> loc(#loc854)
      %31127 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31128 = stablehlo.broadcast_in_dim %31127, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc855)
      %31129 = stablehlo.multiply %31128, %31126 : tensor<1x1x32x2048x2048xbf16> loc(#loc855)
      %31130 = stablehlo.convert %31129 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc856)
      %31131 = func.call @_where_214(%29979, %31130, %iterArg_185) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc857)
      %31132 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc858)
      %31133 = stablehlo.reduce(%31131 init: %31132) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc858)
        stablehlo.return %33956 : tensor<f32> loc(#loc858)
      } loc(#loc858)
      %31134 = stablehlo.broadcast_in_dim %31133, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc859)
      %31135 = stablehlo.broadcast_in_dim %31134, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc860)
      %31136 = stablehlo.subtract %31131, %31135 : tensor<1x1x32x2048x2048xf32> loc(#loc860)
      %31137 = stablehlo.exponential %31136 : tensor<1x1x32x2048x2048xf32> loc(#loc861)
      %31138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc862)
      %31139 = stablehlo.reduce(%31137 init: %31138) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc862)
        stablehlo.return %33956 : tensor<f32> loc(#loc862)
      } loc(#loc862)
      %31140 = stablehlo.broadcast_in_dim %31139, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc859)
      %31141 = stablehlo.broadcast_in_dim %31140, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc863)
      %31142 = stablehlo.divide %31137, %31141 : tensor<1x1x32x2048x2048xf32> loc(#loc863)
      %31143 = stablehlo.convert %31142 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc864)
      %31144 = "stablehlo.slice"(%31107) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc865)
      %31145 = stablehlo.reshape %31144 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc845)
      %31146 = arith.select %ptrue, %31145, %31145 : tensor<1x1x2048x32x64xbf16>
      %31147 = func.call @_einsum_215(%31143, %31146) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc866)
      %31148 = arith.select %ptrue, %31147, %31147 : tensor<1x1x2048x32x64xbf16>
      %31149 = func.call @_where_216(%30042, %iterArg_186, %iterArg_187) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %31150 = func.call @_einsum_217(%31148, %31149) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc867)
      %31151 = arith.select %ptrue, %31150, %31150 : tensor<1x1x2048x2048xbf16>
      %31152 = stablehlo.add %31151, %31067 : tensor<1x1x2048x2048xbf16> loc(#loc869)
      %31153 = stablehlo.convert %31152 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc870)
      %31154 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc871)
      %31155 = stablehlo.reduce(%31153 init: %31154) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc871)
        stablehlo.return %33956 : tensor<f32> loc(#loc871)
      } loc(#loc871)
      %31156 = stablehlo.broadcast_in_dim %31155, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc872)
      %31157 = stablehlo.convert %31156 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc873)
      %31158 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31159 = stablehlo.broadcast_in_dim %31158, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc874)
      %31160 = stablehlo.divide %31157, %31159 : tensor<1x1x2048x1xbf16> loc(#loc874)
      %31161 = stablehlo.broadcast_in_dim %31160, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc875)
      %31162 = stablehlo.subtract %31152, %31161 : tensor<1x1x2048x2048xbf16> loc(#loc875)
      %31163 = stablehlo.broadcast_in_dim %31160, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc875)
      %31164 = stablehlo.subtract %31152, %31163 : tensor<1x1x2048x2048xbf16> loc(#loc875)
      %31165 = stablehlo.multiply %31164, %31164 : tensor<1x1x2048x2048xbf16> loc(#loc876)
      %31166 = stablehlo.convert %31165 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc870)
      %31167 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc871)
      %31168 = stablehlo.reduce(%31166 init: %31167) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc871)
        stablehlo.return %33956 : tensor<f32> loc(#loc871)
      } loc(#loc871)
      %31169 = stablehlo.broadcast_in_dim %31168, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc872)
      %31170 = stablehlo.convert %31169 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc873)
      %31171 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31172 = stablehlo.broadcast_in_dim %31171, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc874)
      %31173 = stablehlo.divide %31170, %31172 : tensor<1x1x2048x1xbf16> loc(#loc874)
      %31174 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31175 = stablehlo.broadcast_in_dim %31174, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc877)
      %31176 = stablehlo.add %31173, %31175 : tensor<1x1x2048x1xbf16> loc(#loc877)
      %31177 = stablehlo.rsqrt %31176 : tensor<1x1x2048x1xbf16> loc(#loc878)
      %31178 = stablehlo.broadcast_in_dim %31177, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc879)
      %31179 = stablehlo.multiply %31162, %31178 : tensor<1x1x2048x2048xbf16> loc(#loc879)
      %31180 = func.call @_where_218(%30042, %iterArg_188, %iterArg_189) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31181 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31182 = stablehlo.broadcast_in_dim %31181, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc880)
      %31183 = stablehlo.add %31182, %31180 : tensor<1x2048xbf16> loc(#loc880)
      %31184 = stablehlo.broadcast_in_dim %31183, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc881)
      %31185 = stablehlo.broadcast_in_dim %31184, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc882)
      %31186 = stablehlo.multiply %31179, %31185 : tensor<1x1x2048x2048xbf16> loc(#loc882)
      %31187 = func.call @_where_219(%30042, %iterArg_190, %iterArg_191) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31188 = stablehlo.broadcast_in_dim %31187, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc883)
      %31189 = stablehlo.broadcast_in_dim %31188, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc884)
      %31190 = stablehlo.add %31186, %31189 : tensor<1x1x2048x2048xbf16> loc(#loc884)
      %31191 = func.call @_where_220(%30042, %iterArg_192, %iterArg_193) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %31192 = func.call @_einsum_221(%31190, %31191) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc885)
      %31193 = arith.select %ptrue, %31192, %31192 : tensor<1x1x2048x8192xbf16>
      %31194 = func.call @_where_222(%30042, %iterArg_194, %iterArg_195) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %31195 = stablehlo.broadcast_in_dim %31194, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc887)
      %31196 = stablehlo.broadcast_in_dim %31195, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc888)
      %31197 = stablehlo.add %31193, %31196 : tensor<1x1x2048x8192xbf16> loc(#loc888)
      %31198 = stablehlo.multiply %31197, %31197 : tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31199 = stablehlo.multiply %31197, %31198 : tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31200 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %31201 = stablehlo.broadcast_in_dim %31200, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31202 = stablehlo.multiply %31201, %31199 : tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31203 = stablehlo.add %31197, %31202 : tensor<1x1x2048x8192xbf16> loc(#loc890)
      %31204 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %31205 = stablehlo.broadcast_in_dim %31204, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31206 = stablehlo.multiply %31205, %31203 : tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31207 = stablehlo.tanh %31206 : tensor<1x1x2048x8192xbf16> loc(#loc891)
      %31208 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31209 = stablehlo.broadcast_in_dim %31208, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc890)
      %31210 = stablehlo.add %31209, %31207 : tensor<1x1x2048x8192xbf16> loc(#loc890)
      %31211 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %31212 = stablehlo.broadcast_in_dim %31211, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31213 = stablehlo.multiply %31212, %31210 : tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31214 = stablehlo.multiply %31197, %31213 : tensor<1x1x2048x8192xbf16> loc(#loc889)
      %31215 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc892)
      %31216 = stablehlo.convert %31215 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc893)
      %31217 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31218 = stablehlo.broadcast_in_dim %31217, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc894)
      %31219 = stablehlo.subtract %31218, %31216 : tensor<1x1x2048x1xf32> loc(#loc894)
      %31220 = stablehlo.convert %31219 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc895)
      %31221 = stablehlo.broadcast_in_dim %31220, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc896)
      %31222 = stablehlo.multiply %31214, %31221 : tensor<1x1x2048x8192xbf16> loc(#loc896)
      %31223 = func.call @_where_223(%30042, %iterArg_196, %iterArg_197) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %31224 = func.call @_einsum_224(%31222, %31223) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc897)
      %31225 = arith.select %ptrue, %31224, %31224 : tensor<1x1x2048x2048xbf16>
      %31226 = func.call @_where_225(%30042, %iterArg_198, %iterArg_199) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31227 = stablehlo.broadcast_in_dim %31226, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc899)
      %31228 = stablehlo.broadcast_in_dim %31227, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc900)
      %31229 = stablehlo.add %31225, %31228 : tensor<1x1x2048x2048xbf16> loc(#loc900)
      %31230 = stablehlo.convert %31215 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc893)
      %31231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31232 = stablehlo.broadcast_in_dim %31231, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc894)
      %31233 = stablehlo.subtract %31232, %31230 : tensor<1x1x2048x1xf32> loc(#loc894)
      %31234 = stablehlo.convert %31233 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc901)
      %31235 = stablehlo.broadcast_in_dim %31234, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc902)
      %31236 = stablehlo.multiply %31229, %31235 : tensor<1x1x2048x2048xbf16> loc(#loc902)
      %31237 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31238 = stablehlo.broadcast_in_dim %31237, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc903)
      %31239 = stablehlo.multiply %31236, %31238 : tensor<1x1x2048x2048xbf16> loc(#loc903)
      %31240 = stablehlo.add %31152, %31239 : tensor<1x1x2048x2048xbf16> loc(#loc904)
      %31241 = stablehlo.convert %31240 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc905)
      %31242 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc906)
      %31243 = stablehlo.reduce(%31241 init: %31242) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc906)
        stablehlo.return %33956 : tensor<f32> loc(#loc906)
      } loc(#loc906)
      %31244 = stablehlo.broadcast_in_dim %31243, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc907)
      %31245 = stablehlo.convert %31244 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc908)
      %31246 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31247 = stablehlo.broadcast_in_dim %31246, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc909)
      %31248 = stablehlo.divide %31245, %31247 : tensor<1x1x2048x1xbf16> loc(#loc909)
      %31249 = stablehlo.broadcast_in_dim %31248, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc910)
      %31250 = stablehlo.subtract %31240, %31249 : tensor<1x1x2048x2048xbf16> loc(#loc910)
      %31251 = stablehlo.broadcast_in_dim %31248, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc910)
      %31252 = stablehlo.subtract %31240, %31251 : tensor<1x1x2048x2048xbf16> loc(#loc910)
      %31253 = stablehlo.multiply %31252, %31252 : tensor<1x1x2048x2048xbf16> loc(#loc911)
      %31254 = stablehlo.convert %31253 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc905)
      %31255 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc906)
      %31256 = stablehlo.reduce(%31254 init: %31255) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc906)
        stablehlo.return %33956 : tensor<f32> loc(#loc906)
      } loc(#loc906)
      %31257 = stablehlo.broadcast_in_dim %31256, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc907)
      %31258 = stablehlo.convert %31257 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc908)
      %31259 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31260 = stablehlo.broadcast_in_dim %31259, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc909)
      %31261 = stablehlo.divide %31258, %31260 : tensor<1x1x2048x1xbf16> loc(#loc909)
      %31262 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31263 = stablehlo.broadcast_in_dim %31262, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc912)
      %31264 = stablehlo.add %31261, %31263 : tensor<1x1x2048x1xbf16> loc(#loc912)
      %31265 = stablehlo.rsqrt %31264 : tensor<1x1x2048x1xbf16> loc(#loc913)
      %31266 = stablehlo.broadcast_in_dim %31265, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc914)
      %31267 = stablehlo.multiply %31250, %31266 : tensor<1x1x2048x2048xbf16> loc(#loc914)
      %31268 = func.call @_where_226(%30042, %iterArg_200, %iterArg_201) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31269 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31270 = stablehlo.broadcast_in_dim %31269, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc915)
      %31271 = stablehlo.add %31270, %31268 : tensor<1x2048xbf16> loc(#loc915)
      %31272 = stablehlo.broadcast_in_dim %31271, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc916)
      %31273 = stablehlo.broadcast_in_dim %31272, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc917)
      %31274 = stablehlo.multiply %31267, %31273 : tensor<1x1x2048x2048xbf16> loc(#loc917)
      %31275 = func.call @_where_227(%30042, %iterArg_202, %iterArg_203) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31276 = stablehlo.broadcast_in_dim %31275, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc918)
      %31277 = stablehlo.broadcast_in_dim %31276, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc919)
      %31278 = stablehlo.add %31274, %31277 : tensor<1x1x2048x2048xbf16> loc(#loc919)
      %31279 = func.call @_where_228(%30042, %iterArg_204, %iterArg_205) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %31280 = func.call @_einsum_229(%31278, %31279) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc920)
      %31281 = "stablehlo.slice"(%31280) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc921)
      %31282 = stablehlo.reshape %31281 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc922)
      %31283 = arith.select %ptrue, %31282, %31282 : tensor<1x1x2048x32x64xbf16>
      %31284 = func.call @_where_230(%30042, %iterArg_206, %iterArg_207) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %31285 = func.call @softplus_231(%31284, %iterArg_208) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc924)
      %31286 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %31287 = stablehlo.broadcast_in_dim %31286, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc925)
      %31288 = stablehlo.multiply %31287, %31285 : tensor<1x64xbf16> loc(#loc925)
      %31289 = stablehlo.broadcast_in_dim %31288, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc926)
      %31290 = stablehlo.broadcast_in_dim %31289, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc927)
      %31291 = stablehlo.multiply %31283, %31290 : tensor<1x1x2048x32x64xbf16> loc(#loc927)
      %31292 = "stablehlo.slice"(%31280) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc928)
      %31293 = stablehlo.reshape %31292 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc922)
      %31294 = arith.select %ptrue, %31293, %31293 : tensor<1x1x2048x32x64xbf16>
      %31295 = func.call @_einsum_233(%31291, %31294) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc929)
      %31296 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31297 = stablehlo.broadcast_in_dim %31296, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc930)
      %31298 = stablehlo.divide %31295, %31297 : tensor<1x1x32x2048x2048xbf16> loc(#loc930)
      %31299 = stablehlo.tanh %31298 : tensor<1x1x32x2048x2048xbf16> loc(#loc931)
      %31300 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31301 = stablehlo.broadcast_in_dim %31300, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc932)
      %31302 = stablehlo.multiply %31301, %31299 : tensor<1x1x32x2048x2048xbf16> loc(#loc932)
      %31303 = stablehlo.convert %31302 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc933)
      %31304 = func.call @_where_234(%29976, %31303, %iterArg_209) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc934)
      %31305 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc935)
      %31306 = stablehlo.reduce(%31304 init: %31305) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc935)
        stablehlo.return %33956 : tensor<f32> loc(#loc935)
      } loc(#loc935)
      %31307 = stablehlo.broadcast_in_dim %31306, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc936)
      %31308 = stablehlo.broadcast_in_dim %31307, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc937)
      %31309 = stablehlo.subtract %31304, %31308 : tensor<1x1x32x2048x2048xf32> loc(#loc937)
      %31310 = stablehlo.exponential %31309 : tensor<1x1x32x2048x2048xf32> loc(#loc938)
      %31311 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc939)
      %31312 = stablehlo.reduce(%31310 init: %31311) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc939)
        stablehlo.return %33956 : tensor<f32> loc(#loc939)
      } loc(#loc939)
      %31313 = stablehlo.broadcast_in_dim %31312, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc936)
      %31314 = stablehlo.broadcast_in_dim %31313, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc940)
      %31315 = stablehlo.divide %31310, %31314 : tensor<1x1x32x2048x2048xf32> loc(#loc940)
      %31316 = stablehlo.convert %31315 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc941)
      %31317 = "stablehlo.slice"(%31280) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc942)
      %31318 = stablehlo.reshape %31317 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc922)
      %31319 = arith.select %ptrue, %31318, %31318 : tensor<1x1x2048x32x64xbf16>
      %31320 = func.call @_einsum_235(%31316, %31319) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc943)
      %31321 = arith.select %ptrue, %31320, %31320 : tensor<1x1x2048x32x64xbf16>
      %31322 = func.call @_where_236(%30042, %iterArg_210, %iterArg_211) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %31323 = func.call @_einsum_237(%31321, %31322) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc944)
      %31324 = arith.select %ptrue, %31323, %31323 : tensor<1x1x2048x2048xbf16>
      %31325 = stablehlo.add %31324, %31240 : tensor<1x1x2048x2048xbf16> loc(#loc946)
      %31326 = stablehlo.convert %31325 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc947)
      %31327 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc948)
      %31328 = stablehlo.reduce(%31326 init: %31327) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc948)
        stablehlo.return %33956 : tensor<f32> loc(#loc948)
      } loc(#loc948)
      %31329 = stablehlo.broadcast_in_dim %31328, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc949)
      %31330 = stablehlo.convert %31329 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc950)
      %31331 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31332 = stablehlo.broadcast_in_dim %31331, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc951)
      %31333 = stablehlo.divide %31330, %31332 : tensor<1x1x2048x1xbf16> loc(#loc951)
      %31334 = stablehlo.broadcast_in_dim %31333, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc952)
      %31335 = stablehlo.subtract %31325, %31334 : tensor<1x1x2048x2048xbf16> loc(#loc952)
      %31336 = stablehlo.broadcast_in_dim %31333, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc952)
      %31337 = stablehlo.subtract %31325, %31336 : tensor<1x1x2048x2048xbf16> loc(#loc952)
      %31338 = stablehlo.multiply %31337, %31337 : tensor<1x1x2048x2048xbf16> loc(#loc953)
      %31339 = stablehlo.convert %31338 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc947)
      %31340 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc948)
      %31341 = stablehlo.reduce(%31339 init: %31340) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc948)
        stablehlo.return %33956 : tensor<f32> loc(#loc948)
      } loc(#loc948)
      %31342 = stablehlo.broadcast_in_dim %31341, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc949)
      %31343 = stablehlo.convert %31342 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc950)
      %31344 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31345 = stablehlo.broadcast_in_dim %31344, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc951)
      %31346 = stablehlo.divide %31343, %31345 : tensor<1x1x2048x1xbf16> loc(#loc951)
      %31347 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31348 = stablehlo.broadcast_in_dim %31347, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc954)
      %31349 = stablehlo.add %31346, %31348 : tensor<1x1x2048x1xbf16> loc(#loc954)
      %31350 = stablehlo.rsqrt %31349 : tensor<1x1x2048x1xbf16> loc(#loc955)
      %31351 = stablehlo.broadcast_in_dim %31350, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc956)
      %31352 = stablehlo.multiply %31335, %31351 : tensor<1x1x2048x2048xbf16> loc(#loc956)
      %31353 = func.call @_where_238(%30042, %iterArg_212, %iterArg_213) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31354 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31355 = stablehlo.broadcast_in_dim %31354, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc957)
      %31356 = stablehlo.add %31355, %31353 : tensor<1x2048xbf16> loc(#loc957)
      %31357 = stablehlo.broadcast_in_dim %31356, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc958)
      %31358 = stablehlo.broadcast_in_dim %31357, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc959)
      %31359 = stablehlo.multiply %31352, %31358 : tensor<1x1x2048x2048xbf16> loc(#loc959)
      %31360 = func.call @_where_239(%30042, %iterArg_214, %iterArg_215) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31361 = stablehlo.broadcast_in_dim %31360, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc960)
      %31362 = stablehlo.broadcast_in_dim %31361, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc961)
      %31363 = stablehlo.add %31359, %31362 : tensor<1x1x2048x2048xbf16> loc(#loc961)
      %31364 = func.call @_where_240(%30042, %iterArg_216, %iterArg_217) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %31365 = func.call @_einsum_241(%31363, %31364) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc962)
      %31366 = arith.select %ptrue, %31365, %31365 : tensor<1x1x2048x8192xbf16>
      %31367 = func.call @_where_242(%30042, %iterArg_218, %iterArg_219) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %31368 = stablehlo.broadcast_in_dim %31367, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc964)
      %31369 = stablehlo.broadcast_in_dim %31368, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc965)
      %31370 = stablehlo.add %31366, %31369 : tensor<1x1x2048x8192xbf16> loc(#loc965)
      %31371 = stablehlo.multiply %31370, %31370 : tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31372 = stablehlo.multiply %31370, %31371 : tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31373 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %31374 = stablehlo.broadcast_in_dim %31373, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31375 = stablehlo.multiply %31374, %31372 : tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31376 = stablehlo.add %31370, %31375 : tensor<1x1x2048x8192xbf16> loc(#loc967)
      %31377 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %31378 = stablehlo.broadcast_in_dim %31377, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31379 = stablehlo.multiply %31378, %31376 : tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31380 = stablehlo.tanh %31379 : tensor<1x1x2048x8192xbf16> loc(#loc968)
      %31381 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31382 = stablehlo.broadcast_in_dim %31381, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc967)
      %31383 = stablehlo.add %31382, %31380 : tensor<1x1x2048x8192xbf16> loc(#loc967)
      %31384 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %31385 = stablehlo.broadcast_in_dim %31384, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31386 = stablehlo.multiply %31385, %31383 : tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31387 = stablehlo.multiply %31370, %31386 : tensor<1x1x2048x8192xbf16> loc(#loc966)
      %31388 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc969)
      %31389 = stablehlo.convert %31388 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc970)
      %31390 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31391 = stablehlo.broadcast_in_dim %31390, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc971)
      %31392 = stablehlo.subtract %31391, %31389 : tensor<1x1x2048x1xf32> loc(#loc971)
      %31393 = stablehlo.convert %31392 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc972)
      %31394 = stablehlo.broadcast_in_dim %31393, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc973)
      %31395 = stablehlo.multiply %31387, %31394 : tensor<1x1x2048x8192xbf16> loc(#loc973)
      %31396 = func.call @_where_243(%30042, %iterArg_220, %iterArg_221) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %31397 = func.call @_einsum_244(%31395, %31396) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc974)
      %31398 = arith.select %ptrue, %31397, %31397 : tensor<1x1x2048x2048xbf16>
      %31399 = func.call @_where_245(%30042, %iterArg_222, %iterArg_223) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31400 = stablehlo.broadcast_in_dim %31399, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc976)
      %31401 = stablehlo.broadcast_in_dim %31400, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc977)
      %31402 = stablehlo.add %31398, %31401 : tensor<1x1x2048x2048xbf16> loc(#loc977)
      %31403 = stablehlo.convert %31388 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc970)
      %31404 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31405 = stablehlo.broadcast_in_dim %31404, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc971)
      %31406 = stablehlo.subtract %31405, %31403 : tensor<1x1x2048x1xf32> loc(#loc971)
      %31407 = stablehlo.convert %31406 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc978)
      %31408 = stablehlo.broadcast_in_dim %31407, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc979)
      %31409 = stablehlo.multiply %31402, %31408 : tensor<1x1x2048x2048xbf16> loc(#loc979)
      %31410 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31411 = stablehlo.broadcast_in_dim %31410, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc980)
      %31412 = stablehlo.multiply %31409, %31411 : tensor<1x1x2048x2048xbf16> loc(#loc980)
      %31413 = stablehlo.add %31325, %31412 : tensor<1x1x2048x2048xbf16> loc(#loc981)
      %31414 = stablehlo.convert %31413 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc982)
      %31415 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc983)
      %31416 = stablehlo.reduce(%31414 init: %31415) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc983)
        stablehlo.return %33956 : tensor<f32> loc(#loc983)
      } loc(#loc983)
      %31417 = stablehlo.broadcast_in_dim %31416, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc984)
      %31418 = stablehlo.convert %31417 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc985)
      %31419 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31420 = stablehlo.broadcast_in_dim %31419, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc986)
      %31421 = stablehlo.divide %31418, %31420 : tensor<1x1x2048x1xbf16> loc(#loc986)
      %31422 = stablehlo.broadcast_in_dim %31421, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc987)
      %31423 = stablehlo.subtract %31413, %31422 : tensor<1x1x2048x2048xbf16> loc(#loc987)
      %31424 = stablehlo.broadcast_in_dim %31421, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc987)
      %31425 = stablehlo.subtract %31413, %31424 : tensor<1x1x2048x2048xbf16> loc(#loc987)
      %31426 = stablehlo.multiply %31425, %31425 : tensor<1x1x2048x2048xbf16> loc(#loc988)
      %31427 = stablehlo.convert %31426 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc982)
      %31428 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc983)
      %31429 = stablehlo.reduce(%31427 init: %31428) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc983)
        stablehlo.return %33956 : tensor<f32> loc(#loc983)
      } loc(#loc983)
      %31430 = stablehlo.broadcast_in_dim %31429, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc984)
      %31431 = stablehlo.convert %31430 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc985)
      %31432 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31433 = stablehlo.broadcast_in_dim %31432, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc986)
      %31434 = stablehlo.divide %31431, %31433 : tensor<1x1x2048x1xbf16> loc(#loc986)
      %31435 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31436 = stablehlo.broadcast_in_dim %31435, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc989)
      %31437 = stablehlo.add %31434, %31436 : tensor<1x1x2048x1xbf16> loc(#loc989)
      %31438 = stablehlo.rsqrt %31437 : tensor<1x1x2048x1xbf16> loc(#loc990)
      %31439 = stablehlo.broadcast_in_dim %31438, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc991)
      %31440 = stablehlo.multiply %31423, %31439 : tensor<1x1x2048x2048xbf16> loc(#loc991)
      %31441 = func.call @_where_246(%30042, %iterArg_224, %iterArg_225) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31442 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31443 = stablehlo.broadcast_in_dim %31442, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc992)
      %31444 = stablehlo.add %31443, %31441 : tensor<1x2048xbf16> loc(#loc992)
      %31445 = stablehlo.broadcast_in_dim %31444, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc993)
      %31446 = stablehlo.broadcast_in_dim %31445, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc994)
      %31447 = stablehlo.multiply %31440, %31446 : tensor<1x1x2048x2048xbf16> loc(#loc994)
      %31448 = func.call @_where_247(%30042, %iterArg_226, %iterArg_227) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31449 = stablehlo.broadcast_in_dim %31448, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc995)
      %31450 = stablehlo.broadcast_in_dim %31449, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc996)
      %31451 = stablehlo.add %31447, %31450 : tensor<1x1x2048x2048xbf16> loc(#loc996)
      %31452 = func.call @_where_248(%30042, %iterArg_228, %iterArg_229) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %31453 = func.call @_einsum_249(%31451, %31452) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc997)
      %31454 = "stablehlo.slice"(%31453) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc998)
      %31455 = stablehlo.reshape %31454 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc999)
      %31456 = arith.select %ptrue, %31455, %31455 : tensor<1x1x2048x32x64xbf16>
      %31457 = func.call @_where_250(%30042, %iterArg_230, %iterArg_231) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %31458 = func.call @softplus_251(%31457, %iterArg_232) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1001)
      %31459 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %31460 = stablehlo.broadcast_in_dim %31459, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1002)
      %31461 = stablehlo.multiply %31460, %31458 : tensor<1x64xbf16> loc(#loc1002)
      %31462 = stablehlo.broadcast_in_dim %31461, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1003)
      %31463 = stablehlo.broadcast_in_dim %31462, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1004)
      %31464 = stablehlo.multiply %31456, %31463 : tensor<1x1x2048x32x64xbf16> loc(#loc1004)
      %31465 = "stablehlo.slice"(%31453) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1005)
      %31466 = stablehlo.reshape %31465 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc999)
      %31467 = arith.select %ptrue, %31466, %31466 : tensor<1x1x2048x32x64xbf16>
      %31468 = func.call @_einsum_253(%31464, %31467) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1006)
      %31469 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31470 = stablehlo.broadcast_in_dim %31469, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1007)
      %31471 = stablehlo.divide %31468, %31470 : tensor<1x1x32x2048x2048xbf16> loc(#loc1007)
      %31472 = stablehlo.tanh %31471 : tensor<1x1x32x2048x2048xbf16> loc(#loc1008)
      %31473 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31474 = stablehlo.broadcast_in_dim %31473, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1009)
      %31475 = stablehlo.multiply %31474, %31472 : tensor<1x1x32x2048x2048xbf16> loc(#loc1009)
      %31476 = stablehlo.convert %31475 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1010)
      %31477 = func.call @_where_254(%29973, %31476, %iterArg_233) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1011)
      %31478 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1012)
      %31479 = stablehlo.reduce(%31477 init: %31478) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1012)
        stablehlo.return %33956 : tensor<f32> loc(#loc1012)
      } loc(#loc1012)
      %31480 = stablehlo.broadcast_in_dim %31479, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1013)
      %31481 = stablehlo.broadcast_in_dim %31480, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1014)
      %31482 = stablehlo.subtract %31477, %31481 : tensor<1x1x32x2048x2048xf32> loc(#loc1014)
      %31483 = stablehlo.exponential %31482 : tensor<1x1x32x2048x2048xf32> loc(#loc1015)
      %31484 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1016)
      %31485 = stablehlo.reduce(%31483 init: %31484) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1016)
        stablehlo.return %33956 : tensor<f32> loc(#loc1016)
      } loc(#loc1016)
      %31486 = stablehlo.broadcast_in_dim %31485, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1013)
      %31487 = stablehlo.broadcast_in_dim %31486, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1017)
      %31488 = stablehlo.divide %31483, %31487 : tensor<1x1x32x2048x2048xf32> loc(#loc1017)
      %31489 = stablehlo.convert %31488 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1018)
      %31490 = "stablehlo.slice"(%31453) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1019)
      %31491 = stablehlo.reshape %31490 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc999)
      %31492 = arith.select %ptrue, %31491, %31491 : tensor<1x1x2048x32x64xbf16>
      %31493 = func.call @_einsum_255(%31489, %31492) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1020)
      %31494 = arith.select %ptrue, %31493, %31493 : tensor<1x1x2048x32x64xbf16>
      %31495 = func.call @_where_256(%30042, %iterArg_234, %iterArg_235) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %31496 = func.call @_einsum_257(%31494, %31495) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1021)
      %31497 = arith.select %ptrue, %31496, %31496 : tensor<1x1x2048x2048xbf16>
      %31498 = stablehlo.add %31497, %31413 : tensor<1x1x2048x2048xbf16> loc(#loc1023)
      %31499 = stablehlo.convert %31498 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1024)
      %31500 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1025)
      %31501 = stablehlo.reduce(%31499 init: %31500) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1025)
        stablehlo.return %33956 : tensor<f32> loc(#loc1025)
      } loc(#loc1025)
      %31502 = stablehlo.broadcast_in_dim %31501, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1026)
      %31503 = stablehlo.convert %31502 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1027)
      %31504 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31505 = stablehlo.broadcast_in_dim %31504, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1028)
      %31506 = stablehlo.divide %31503, %31505 : tensor<1x1x2048x1xbf16> loc(#loc1028)
      %31507 = stablehlo.broadcast_in_dim %31506, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1029)
      %31508 = stablehlo.subtract %31498, %31507 : tensor<1x1x2048x2048xbf16> loc(#loc1029)
      %31509 = stablehlo.broadcast_in_dim %31506, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1029)
      %31510 = stablehlo.subtract %31498, %31509 : tensor<1x1x2048x2048xbf16> loc(#loc1029)
      %31511 = stablehlo.multiply %31510, %31510 : tensor<1x1x2048x2048xbf16> loc(#loc1030)
      %31512 = stablehlo.convert %31511 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1024)
      %31513 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1025)
      %31514 = stablehlo.reduce(%31512 init: %31513) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1025)
        stablehlo.return %33956 : tensor<f32> loc(#loc1025)
      } loc(#loc1025)
      %31515 = stablehlo.broadcast_in_dim %31514, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1026)
      %31516 = stablehlo.convert %31515 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1027)
      %31517 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31518 = stablehlo.broadcast_in_dim %31517, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1028)
      %31519 = stablehlo.divide %31516, %31518 : tensor<1x1x2048x1xbf16> loc(#loc1028)
      %31520 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31521 = stablehlo.broadcast_in_dim %31520, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1031)
      %31522 = stablehlo.add %31519, %31521 : tensor<1x1x2048x1xbf16> loc(#loc1031)
      %31523 = stablehlo.rsqrt %31522 : tensor<1x1x2048x1xbf16> loc(#loc1032)
      %31524 = stablehlo.broadcast_in_dim %31523, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1033)
      %31525 = stablehlo.multiply %31508, %31524 : tensor<1x1x2048x2048xbf16> loc(#loc1033)
      %31526 = func.call @_where_258(%30042, %iterArg_236, %iterArg_237) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31527 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31528 = stablehlo.broadcast_in_dim %31527, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1034)
      %31529 = stablehlo.add %31528, %31526 : tensor<1x2048xbf16> loc(#loc1034)
      %31530 = stablehlo.broadcast_in_dim %31529, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1035)
      %31531 = stablehlo.broadcast_in_dim %31530, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1036)
      %31532 = stablehlo.multiply %31525, %31531 : tensor<1x1x2048x2048xbf16> loc(#loc1036)
      %31533 = func.call @_where_259(%30042, %iterArg_238, %iterArg_239) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31534 = stablehlo.broadcast_in_dim %31533, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1037)
      %31535 = stablehlo.broadcast_in_dim %31534, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1038)
      %31536 = stablehlo.add %31532, %31535 : tensor<1x1x2048x2048xbf16> loc(#loc1038)
      %31537 = func.call @_where_260(%30042, %iterArg_240, %iterArg_241) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %31538 = func.call @_einsum_261(%31536, %31537) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1039)
      %31539 = arith.select %ptrue, %31538, %31538 : tensor<1x1x2048x8192xbf16>
      %31540 = func.call @_where_262(%30042, %iterArg_242, %iterArg_243) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %31541 = stablehlo.broadcast_in_dim %31540, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1041)
      %31542 = stablehlo.broadcast_in_dim %31541, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1042)
      %31543 = stablehlo.add %31539, %31542 : tensor<1x1x2048x8192xbf16> loc(#loc1042)
      %31544 = stablehlo.multiply %31543, %31543 : tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31545 = stablehlo.multiply %31543, %31544 : tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31546 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %31547 = stablehlo.broadcast_in_dim %31546, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31548 = stablehlo.multiply %31547, %31545 : tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31549 = stablehlo.add %31543, %31548 : tensor<1x1x2048x8192xbf16> loc(#loc1044)
      %31550 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %31551 = stablehlo.broadcast_in_dim %31550, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31552 = stablehlo.multiply %31551, %31549 : tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31553 = stablehlo.tanh %31552 : tensor<1x1x2048x8192xbf16> loc(#loc1045)
      %31554 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31555 = stablehlo.broadcast_in_dim %31554, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1044)
      %31556 = stablehlo.add %31555, %31553 : tensor<1x1x2048x8192xbf16> loc(#loc1044)
      %31557 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %31558 = stablehlo.broadcast_in_dim %31557, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31559 = stablehlo.multiply %31558, %31556 : tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31560 = stablehlo.multiply %31543, %31559 : tensor<1x1x2048x8192xbf16> loc(#loc1043)
      %31561 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1046)
      %31562 = stablehlo.convert %31561 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1047)
      %31563 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31564 = stablehlo.broadcast_in_dim %31563, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1048)
      %31565 = stablehlo.subtract %31564, %31562 : tensor<1x1x2048x1xf32> loc(#loc1048)
      %31566 = stablehlo.convert %31565 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1049)
      %31567 = stablehlo.broadcast_in_dim %31566, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1050)
      %31568 = stablehlo.multiply %31560, %31567 : tensor<1x1x2048x8192xbf16> loc(#loc1050)
      %31569 = func.call @_where_263(%30042, %iterArg_244, %iterArg_245) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %31570 = func.call @_einsum_264(%31568, %31569) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1051)
      %31571 = arith.select %ptrue, %31570, %31570 : tensor<1x1x2048x2048xbf16>
      %31572 = func.call @_where_265(%30042, %iterArg_246, %iterArg_247) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31573 = stablehlo.broadcast_in_dim %31572, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1053)
      %31574 = stablehlo.broadcast_in_dim %31573, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1054)
      %31575 = stablehlo.add %31571, %31574 : tensor<1x1x2048x2048xbf16> loc(#loc1054)
      %31576 = stablehlo.convert %31561 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1047)
      %31577 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31578 = stablehlo.broadcast_in_dim %31577, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1048)
      %31579 = stablehlo.subtract %31578, %31576 : tensor<1x1x2048x1xf32> loc(#loc1048)
      %31580 = stablehlo.convert %31579 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1055)
      %31581 = stablehlo.broadcast_in_dim %31580, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1056)
      %31582 = stablehlo.multiply %31575, %31581 : tensor<1x1x2048x2048xbf16> loc(#loc1056)
      %31583 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31584 = stablehlo.broadcast_in_dim %31583, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1057)
      %31585 = stablehlo.multiply %31582, %31584 : tensor<1x1x2048x2048xbf16> loc(#loc1057)
      %31586 = stablehlo.add %31498, %31585 : tensor<1x1x2048x2048xbf16> loc(#loc1058)
      %31587 = stablehlo.convert %31586 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1059)
      %31588 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1060)
      %31589 = stablehlo.reduce(%31587 init: %31588) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1060)
        stablehlo.return %33956 : tensor<f32> loc(#loc1060)
      } loc(#loc1060)
      %31590 = stablehlo.broadcast_in_dim %31589, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1061)
      %31591 = stablehlo.convert %31590 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1062)
      %31592 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31593 = stablehlo.broadcast_in_dim %31592, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1063)
      %31594 = stablehlo.divide %31591, %31593 : tensor<1x1x2048x1xbf16> loc(#loc1063)
      %31595 = stablehlo.broadcast_in_dim %31594, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1064)
      %31596 = stablehlo.subtract %31586, %31595 : tensor<1x1x2048x2048xbf16> loc(#loc1064)
      %31597 = stablehlo.broadcast_in_dim %31594, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1064)
      %31598 = stablehlo.subtract %31586, %31597 : tensor<1x1x2048x2048xbf16> loc(#loc1064)
      %31599 = stablehlo.multiply %31598, %31598 : tensor<1x1x2048x2048xbf16> loc(#loc1065)
      %31600 = stablehlo.convert %31599 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1059)
      %31601 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1060)
      %31602 = stablehlo.reduce(%31600 init: %31601) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1060)
        stablehlo.return %33956 : tensor<f32> loc(#loc1060)
      } loc(#loc1060)
      %31603 = stablehlo.broadcast_in_dim %31602, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1061)
      %31604 = stablehlo.convert %31603 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1062)
      %31605 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31606 = stablehlo.broadcast_in_dim %31605, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1063)
      %31607 = stablehlo.divide %31604, %31606 : tensor<1x1x2048x1xbf16> loc(#loc1063)
      %31608 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31609 = stablehlo.broadcast_in_dim %31608, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1066)
      %31610 = stablehlo.add %31607, %31609 : tensor<1x1x2048x1xbf16> loc(#loc1066)
      %31611 = stablehlo.rsqrt %31610 : tensor<1x1x2048x1xbf16> loc(#loc1067)
      %31612 = stablehlo.broadcast_in_dim %31611, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1068)
      %31613 = stablehlo.multiply %31596, %31612 : tensor<1x1x2048x2048xbf16> loc(#loc1068)
      %31614 = func.call @_where_266(%30042, %iterArg_248, %iterArg_249) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31615 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31616 = stablehlo.broadcast_in_dim %31615, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1069)
      %31617 = stablehlo.add %31616, %31614 : tensor<1x2048xbf16> loc(#loc1069)
      %31618 = stablehlo.broadcast_in_dim %31617, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1070)
      %31619 = stablehlo.broadcast_in_dim %31618, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1071)
      %31620 = stablehlo.multiply %31613, %31619 : tensor<1x1x2048x2048xbf16> loc(#loc1071)
      %31621 = func.call @_where_267(%30042, %iterArg_250, %iterArg_251) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31622 = stablehlo.broadcast_in_dim %31621, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1072)
      %31623 = stablehlo.broadcast_in_dim %31622, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1073)
      %31624 = stablehlo.add %31620, %31623 : tensor<1x1x2048x2048xbf16> loc(#loc1073)
      %31625 = func.call @_where_268(%30042, %iterArg_252, %iterArg_253) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %31626 = func.call @_einsum_269(%31624, %31625) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1074)
      %31627 = "stablehlo.slice"(%31626) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1075)
      %31628 = stablehlo.reshape %31627 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1076)
      %31629 = arith.select %ptrue, %31628, %31628 : tensor<1x1x2048x32x64xbf16>
      %31630 = func.call @_where_270(%30042, %iterArg_254, %iterArg_255) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %31631 = func.call @softplus_271(%31630, %iterArg_256) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1078)
      %31632 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %31633 = stablehlo.broadcast_in_dim %31632, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1079)
      %31634 = stablehlo.multiply %31633, %31631 : tensor<1x64xbf16> loc(#loc1079)
      %31635 = stablehlo.broadcast_in_dim %31634, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1080)
      %31636 = stablehlo.broadcast_in_dim %31635, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1081)
      %31637 = stablehlo.multiply %31629, %31636 : tensor<1x1x2048x32x64xbf16> loc(#loc1081)
      %31638 = "stablehlo.slice"(%31626) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1082)
      %31639 = stablehlo.reshape %31638 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1076)
      %31640 = arith.select %ptrue, %31639, %31639 : tensor<1x1x2048x32x64xbf16>
      %31641 = func.call @_einsum_273(%31637, %31640) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1083)
      %31642 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31643 = stablehlo.broadcast_in_dim %31642, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1084)
      %31644 = stablehlo.divide %31641, %31643 : tensor<1x1x32x2048x2048xbf16> loc(#loc1084)
      %31645 = stablehlo.tanh %31644 : tensor<1x1x32x2048x2048xbf16> loc(#loc1085)
      %31646 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31647 = stablehlo.broadcast_in_dim %31646, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1086)
      %31648 = stablehlo.multiply %31647, %31645 : tensor<1x1x32x2048x2048xbf16> loc(#loc1086)
      %31649 = stablehlo.convert %31648 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1087)
      %31650 = func.call @_where_274(%29970, %31649, %iterArg_257) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1088)
      %31651 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1089)
      %31652 = stablehlo.reduce(%31650 init: %31651) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1089)
        stablehlo.return %33956 : tensor<f32> loc(#loc1089)
      } loc(#loc1089)
      %31653 = stablehlo.broadcast_in_dim %31652, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1090)
      %31654 = stablehlo.broadcast_in_dim %31653, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1091)
      %31655 = stablehlo.subtract %31650, %31654 : tensor<1x1x32x2048x2048xf32> loc(#loc1091)
      %31656 = stablehlo.exponential %31655 : tensor<1x1x32x2048x2048xf32> loc(#loc1092)
      %31657 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1093)
      %31658 = stablehlo.reduce(%31656 init: %31657) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1093)
        stablehlo.return %33956 : tensor<f32> loc(#loc1093)
      } loc(#loc1093)
      %31659 = stablehlo.broadcast_in_dim %31658, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1090)
      %31660 = stablehlo.broadcast_in_dim %31659, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1094)
      %31661 = stablehlo.divide %31656, %31660 : tensor<1x1x32x2048x2048xf32> loc(#loc1094)
      %31662 = stablehlo.convert %31661 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1095)
      %31663 = "stablehlo.slice"(%31626) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1096)
      %31664 = stablehlo.reshape %31663 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1076)
      %31665 = arith.select %ptrue, %31664, %31664 : tensor<1x1x2048x32x64xbf16>
      %31666 = func.call @_einsum_275(%31662, %31665) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1097)
      %31667 = arith.select %ptrue, %31666, %31666 : tensor<1x1x2048x32x64xbf16>
      %31668 = func.call @_where_276(%30042, %iterArg_258, %iterArg_259) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %31669 = func.call @_einsum_277(%31667, %31668) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1098)
      %31670 = arith.select %ptrue, %31669, %31669 : tensor<1x1x2048x2048xbf16>
      %31671 = stablehlo.add %31670, %31586 : tensor<1x1x2048x2048xbf16> loc(#loc1100)
      %31672 = stablehlo.convert %31671 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1101)
      %31673 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1102)
      %31674 = stablehlo.reduce(%31672 init: %31673) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1102)
        stablehlo.return %33956 : tensor<f32> loc(#loc1102)
      } loc(#loc1102)
      %31675 = stablehlo.broadcast_in_dim %31674, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1103)
      %31676 = stablehlo.convert %31675 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1104)
      %31677 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31678 = stablehlo.broadcast_in_dim %31677, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1105)
      %31679 = stablehlo.divide %31676, %31678 : tensor<1x1x2048x1xbf16> loc(#loc1105)
      %31680 = stablehlo.broadcast_in_dim %31679, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1106)
      %31681 = stablehlo.subtract %31671, %31680 : tensor<1x1x2048x2048xbf16> loc(#loc1106)
      %31682 = stablehlo.broadcast_in_dim %31679, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1106)
      %31683 = stablehlo.subtract %31671, %31682 : tensor<1x1x2048x2048xbf16> loc(#loc1106)
      %31684 = stablehlo.multiply %31683, %31683 : tensor<1x1x2048x2048xbf16> loc(#loc1107)
      %31685 = stablehlo.convert %31684 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1101)
      %31686 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1102)
      %31687 = stablehlo.reduce(%31685 init: %31686) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1102)
        stablehlo.return %33956 : tensor<f32> loc(#loc1102)
      } loc(#loc1102)
      %31688 = stablehlo.broadcast_in_dim %31687, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1103)
      %31689 = stablehlo.convert %31688 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1104)
      %31690 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31691 = stablehlo.broadcast_in_dim %31690, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1105)
      %31692 = stablehlo.divide %31689, %31691 : tensor<1x1x2048x1xbf16> loc(#loc1105)
      %31693 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31694 = stablehlo.broadcast_in_dim %31693, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1108)
      %31695 = stablehlo.add %31692, %31694 : tensor<1x1x2048x1xbf16> loc(#loc1108)
      %31696 = stablehlo.rsqrt %31695 : tensor<1x1x2048x1xbf16> loc(#loc1109)
      %31697 = stablehlo.broadcast_in_dim %31696, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1110)
      %31698 = stablehlo.multiply %31681, %31697 : tensor<1x1x2048x2048xbf16> loc(#loc1110)
      %31699 = func.call @_where_278(%30042, %iterArg_260, %iterArg_261) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31700 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31701 = stablehlo.broadcast_in_dim %31700, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1111)
      %31702 = stablehlo.add %31701, %31699 : tensor<1x2048xbf16> loc(#loc1111)
      %31703 = stablehlo.broadcast_in_dim %31702, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1112)
      %31704 = stablehlo.broadcast_in_dim %31703, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1113)
      %31705 = stablehlo.multiply %31698, %31704 : tensor<1x1x2048x2048xbf16> loc(#loc1113)
      %31706 = func.call @_where_279(%30042, %iterArg_262, %iterArg_263) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31707 = stablehlo.broadcast_in_dim %31706, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1114)
      %31708 = stablehlo.broadcast_in_dim %31707, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1115)
      %31709 = stablehlo.add %31705, %31708 : tensor<1x1x2048x2048xbf16> loc(#loc1115)
      %31710 = func.call @_where_280(%30042, %iterArg_264, %iterArg_265) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %31711 = func.call @_einsum_281(%31709, %31710) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1116)
      %31712 = arith.select %ptrue, %31711, %31711 : tensor<1x1x2048x8192xbf16>
      %31713 = func.call @_where_282(%30042, %iterArg_266, %iterArg_267) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %31714 = stablehlo.broadcast_in_dim %31713, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1118)
      %31715 = stablehlo.broadcast_in_dim %31714, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1119)
      %31716 = stablehlo.add %31712, %31715 : tensor<1x1x2048x8192xbf16> loc(#loc1119)
      %31717 = stablehlo.multiply %31716, %31716 : tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31718 = stablehlo.multiply %31716, %31717 : tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31719 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %31720 = stablehlo.broadcast_in_dim %31719, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31721 = stablehlo.multiply %31720, %31718 : tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31722 = stablehlo.add %31716, %31721 : tensor<1x1x2048x8192xbf16> loc(#loc1121)
      %31723 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %31724 = stablehlo.broadcast_in_dim %31723, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31725 = stablehlo.multiply %31724, %31722 : tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31726 = stablehlo.tanh %31725 : tensor<1x1x2048x8192xbf16> loc(#loc1122)
      %31727 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31728 = stablehlo.broadcast_in_dim %31727, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1121)
      %31729 = stablehlo.add %31728, %31726 : tensor<1x1x2048x8192xbf16> loc(#loc1121)
      %31730 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %31731 = stablehlo.broadcast_in_dim %31730, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31732 = stablehlo.multiply %31731, %31729 : tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31733 = stablehlo.multiply %31716, %31732 : tensor<1x1x2048x8192xbf16> loc(#loc1120)
      %31734 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1123)
      %31735 = stablehlo.convert %31734 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1124)
      %31736 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31737 = stablehlo.broadcast_in_dim %31736, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1125)
      %31738 = stablehlo.subtract %31737, %31735 : tensor<1x1x2048x1xf32> loc(#loc1125)
      %31739 = stablehlo.convert %31738 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1126)
      %31740 = stablehlo.broadcast_in_dim %31739, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1127)
      %31741 = stablehlo.multiply %31733, %31740 : tensor<1x1x2048x8192xbf16> loc(#loc1127)
      %31742 = func.call @_where_283(%30042, %iterArg_268, %iterArg_269) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %31743 = func.call @_einsum_284(%31741, %31742) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1128)
      %31744 = arith.select %ptrue, %31743, %31743 : tensor<1x1x2048x2048xbf16>
      %31745 = func.call @_where_285(%30042, %iterArg_270, %iterArg_271) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31746 = stablehlo.broadcast_in_dim %31745, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1130)
      %31747 = stablehlo.broadcast_in_dim %31746, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1131)
      %31748 = stablehlo.add %31744, %31747 : tensor<1x1x2048x2048xbf16> loc(#loc1131)
      %31749 = stablehlo.convert %31734 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1124)
      %31750 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31751 = stablehlo.broadcast_in_dim %31750, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1125)
      %31752 = stablehlo.subtract %31751, %31749 : tensor<1x1x2048x1xf32> loc(#loc1125)
      %31753 = stablehlo.convert %31752 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1132)
      %31754 = stablehlo.broadcast_in_dim %31753, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1133)
      %31755 = stablehlo.multiply %31748, %31754 : tensor<1x1x2048x2048xbf16> loc(#loc1133)
      %31756 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31757 = stablehlo.broadcast_in_dim %31756, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1134)
      %31758 = stablehlo.multiply %31755, %31757 : tensor<1x1x2048x2048xbf16> loc(#loc1134)
      %31759 = stablehlo.add %31671, %31758 : tensor<1x1x2048x2048xbf16> loc(#loc1135)
      %31760 = stablehlo.convert %31759 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1136)
      %31761 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1137)
      %31762 = stablehlo.reduce(%31760 init: %31761) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1137)
        stablehlo.return %33956 : tensor<f32> loc(#loc1137)
      } loc(#loc1137)
      %31763 = stablehlo.broadcast_in_dim %31762, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1138)
      %31764 = stablehlo.convert %31763 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1139)
      %31765 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31766 = stablehlo.broadcast_in_dim %31765, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1140)
      %31767 = stablehlo.divide %31764, %31766 : tensor<1x1x2048x1xbf16> loc(#loc1140)
      %31768 = stablehlo.broadcast_in_dim %31767, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1141)
      %31769 = stablehlo.subtract %31759, %31768 : tensor<1x1x2048x2048xbf16> loc(#loc1141)
      %31770 = stablehlo.broadcast_in_dim %31767, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1141)
      %31771 = stablehlo.subtract %31759, %31770 : tensor<1x1x2048x2048xbf16> loc(#loc1141)
      %31772 = stablehlo.multiply %31771, %31771 : tensor<1x1x2048x2048xbf16> loc(#loc1142)
      %31773 = stablehlo.convert %31772 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1136)
      %31774 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1137)
      %31775 = stablehlo.reduce(%31773 init: %31774) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1137)
        stablehlo.return %33956 : tensor<f32> loc(#loc1137)
      } loc(#loc1137)
      %31776 = stablehlo.broadcast_in_dim %31775, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1138)
      %31777 = stablehlo.convert %31776 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1139)
      %31778 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31779 = stablehlo.broadcast_in_dim %31778, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1140)
      %31780 = stablehlo.divide %31777, %31779 : tensor<1x1x2048x1xbf16> loc(#loc1140)
      %31781 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31782 = stablehlo.broadcast_in_dim %31781, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1143)
      %31783 = stablehlo.add %31780, %31782 : tensor<1x1x2048x1xbf16> loc(#loc1143)
      %31784 = stablehlo.rsqrt %31783 : tensor<1x1x2048x1xbf16> loc(#loc1144)
      %31785 = stablehlo.broadcast_in_dim %31784, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1145)
      %31786 = stablehlo.multiply %31769, %31785 : tensor<1x1x2048x2048xbf16> loc(#loc1145)
      %31787 = func.call @_where_286(%30042, %iterArg_272, %iterArg_273) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31788 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31789 = stablehlo.broadcast_in_dim %31788, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1146)
      %31790 = stablehlo.add %31789, %31787 : tensor<1x2048xbf16> loc(#loc1146)
      %31791 = stablehlo.broadcast_in_dim %31790, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1147)
      %31792 = stablehlo.broadcast_in_dim %31791, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1148)
      %31793 = stablehlo.multiply %31786, %31792 : tensor<1x1x2048x2048xbf16> loc(#loc1148)
      %31794 = func.call @_where_287(%30042, %iterArg_274, %iterArg_275) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31795 = stablehlo.broadcast_in_dim %31794, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1149)
      %31796 = stablehlo.broadcast_in_dim %31795, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1150)
      %31797 = stablehlo.add %31793, %31796 : tensor<1x1x2048x2048xbf16> loc(#loc1150)
      %31798 = func.call @_where_288(%30042, %iterArg_276, %iterArg_277) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %31799 = func.call @_einsum_289(%31797, %31798) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1151)
      %31800 = "stablehlo.slice"(%31799) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1152)
      %31801 = stablehlo.reshape %31800 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1153)
      %31802 = arith.select %ptrue, %31801, %31801 : tensor<1x1x2048x32x64xbf16>
      %31803 = func.call @_where_290(%30042, %iterArg_278, %iterArg_279) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %31804 = func.call @softplus_291(%31803, %iterArg_280) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1155)
      %31805 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %31806 = stablehlo.broadcast_in_dim %31805, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1156)
      %31807 = stablehlo.multiply %31806, %31804 : tensor<1x64xbf16> loc(#loc1156)
      %31808 = stablehlo.broadcast_in_dim %31807, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1157)
      %31809 = stablehlo.broadcast_in_dim %31808, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1158)
      %31810 = stablehlo.multiply %31802, %31809 : tensor<1x1x2048x32x64xbf16> loc(#loc1158)
      %31811 = "stablehlo.slice"(%31799) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1159)
      %31812 = stablehlo.reshape %31811 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1153)
      %31813 = arith.select %ptrue, %31812, %31812 : tensor<1x1x2048x32x64xbf16>
      %31814 = func.call @_einsum_293(%31810, %31813) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1160)
      %31815 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31816 = stablehlo.broadcast_in_dim %31815, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1161)
      %31817 = stablehlo.divide %31814, %31816 : tensor<1x1x32x2048x2048xbf16> loc(#loc1161)
      %31818 = stablehlo.tanh %31817 : tensor<1x1x32x2048x2048xbf16> loc(#loc1162)
      %31819 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31820 = stablehlo.broadcast_in_dim %31819, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1163)
      %31821 = stablehlo.multiply %31820, %31818 : tensor<1x1x32x2048x2048xbf16> loc(#loc1163)
      %31822 = stablehlo.convert %31821 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1164)
      %31823 = func.call @_where_294(%29967, %31822, %iterArg_281) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1165)
      %31824 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1166)
      %31825 = stablehlo.reduce(%31823 init: %31824) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1166)
        stablehlo.return %33956 : tensor<f32> loc(#loc1166)
      } loc(#loc1166)
      %31826 = stablehlo.broadcast_in_dim %31825, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1167)
      %31827 = stablehlo.broadcast_in_dim %31826, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1168)
      %31828 = stablehlo.subtract %31823, %31827 : tensor<1x1x32x2048x2048xf32> loc(#loc1168)
      %31829 = stablehlo.exponential %31828 : tensor<1x1x32x2048x2048xf32> loc(#loc1169)
      %31830 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1170)
      %31831 = stablehlo.reduce(%31829 init: %31830) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1170)
        stablehlo.return %33956 : tensor<f32> loc(#loc1170)
      } loc(#loc1170)
      %31832 = stablehlo.broadcast_in_dim %31831, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1167)
      %31833 = stablehlo.broadcast_in_dim %31832, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1171)
      %31834 = stablehlo.divide %31829, %31833 : tensor<1x1x32x2048x2048xf32> loc(#loc1171)
      %31835 = stablehlo.convert %31834 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1172)
      %31836 = "stablehlo.slice"(%31799) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1173)
      %31837 = stablehlo.reshape %31836 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1153)
      %31838 = arith.select %ptrue, %31837, %31837 : tensor<1x1x2048x32x64xbf16>
      %31839 = func.call @_einsum_295(%31835, %31838) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1174)
      %31840 = arith.select %ptrue, %31839, %31839 : tensor<1x1x2048x32x64xbf16>
      %31841 = func.call @_where_296(%30042, %iterArg_282, %iterArg_283) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %31842 = func.call @_einsum_297(%31840, %31841) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1175)
      %31843 = arith.select %ptrue, %31842, %31842 : tensor<1x1x2048x2048xbf16>
      %31844 = stablehlo.add %31843, %31759 : tensor<1x1x2048x2048xbf16> loc(#loc1177)
      %31845 = stablehlo.convert %31844 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1178)
      %31846 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1179)
      %31847 = stablehlo.reduce(%31845 init: %31846) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1179)
        stablehlo.return %33956 : tensor<f32> loc(#loc1179)
      } loc(#loc1179)
      %31848 = stablehlo.broadcast_in_dim %31847, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1180)
      %31849 = stablehlo.convert %31848 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1181)
      %31850 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31851 = stablehlo.broadcast_in_dim %31850, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1182)
      %31852 = stablehlo.divide %31849, %31851 : tensor<1x1x2048x1xbf16> loc(#loc1182)
      %31853 = stablehlo.broadcast_in_dim %31852, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1183)
      %31854 = stablehlo.subtract %31844, %31853 : tensor<1x1x2048x2048xbf16> loc(#loc1183)
      %31855 = stablehlo.broadcast_in_dim %31852, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1183)
      %31856 = stablehlo.subtract %31844, %31855 : tensor<1x1x2048x2048xbf16> loc(#loc1183)
      %31857 = stablehlo.multiply %31856, %31856 : tensor<1x1x2048x2048xbf16> loc(#loc1184)
      %31858 = stablehlo.convert %31857 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1178)
      %31859 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1179)
      %31860 = stablehlo.reduce(%31858 init: %31859) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1179)
        stablehlo.return %33956 : tensor<f32> loc(#loc1179)
      } loc(#loc1179)
      %31861 = stablehlo.broadcast_in_dim %31860, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1180)
      %31862 = stablehlo.convert %31861 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1181)
      %31863 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31864 = stablehlo.broadcast_in_dim %31863, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1182)
      %31865 = stablehlo.divide %31862, %31864 : tensor<1x1x2048x1xbf16> loc(#loc1182)
      %31866 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31867 = stablehlo.broadcast_in_dim %31866, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1185)
      %31868 = stablehlo.add %31865, %31867 : tensor<1x1x2048x1xbf16> loc(#loc1185)
      %31869 = stablehlo.rsqrt %31868 : tensor<1x1x2048x1xbf16> loc(#loc1186)
      %31870 = stablehlo.broadcast_in_dim %31869, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1187)
      %31871 = stablehlo.multiply %31854, %31870 : tensor<1x1x2048x2048xbf16> loc(#loc1187)
      %31872 = func.call @_where_298(%30042, %iterArg_284, %iterArg_285) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31873 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31874 = stablehlo.broadcast_in_dim %31873, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1188)
      %31875 = stablehlo.add %31874, %31872 : tensor<1x2048xbf16> loc(#loc1188)
      %31876 = stablehlo.broadcast_in_dim %31875, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1189)
      %31877 = stablehlo.broadcast_in_dim %31876, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1190)
      %31878 = stablehlo.multiply %31871, %31877 : tensor<1x1x2048x2048xbf16> loc(#loc1190)
      %31879 = func.call @_where_299(%30042, %iterArg_286, %iterArg_287) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31880 = stablehlo.broadcast_in_dim %31879, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1191)
      %31881 = stablehlo.broadcast_in_dim %31880, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1192)
      %31882 = stablehlo.add %31878, %31881 : tensor<1x1x2048x2048xbf16> loc(#loc1192)
      %31883 = func.call @_where_300(%30042, %iterArg_288, %iterArg_289) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %31884 = func.call @_einsum_301(%31882, %31883) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1193)
      %31885 = arith.select %ptrue, %31884, %31884 : tensor<1x1x2048x8192xbf16>
      %31886 = func.call @_where_302(%30042, %iterArg_290, %iterArg_291) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %31887 = stablehlo.broadcast_in_dim %31886, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1195)
      %31888 = stablehlo.broadcast_in_dim %31887, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1196)
      %31889 = stablehlo.add %31885, %31888 : tensor<1x1x2048x8192xbf16> loc(#loc1196)
      %31890 = stablehlo.multiply %31889, %31889 : tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31891 = stablehlo.multiply %31889, %31890 : tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31892 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %31893 = stablehlo.broadcast_in_dim %31892, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31894 = stablehlo.multiply %31893, %31891 : tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31895 = stablehlo.add %31889, %31894 : tensor<1x1x2048x8192xbf16> loc(#loc1198)
      %31896 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %31897 = stablehlo.broadcast_in_dim %31896, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31898 = stablehlo.multiply %31897, %31895 : tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31899 = stablehlo.tanh %31898 : tensor<1x1x2048x8192xbf16> loc(#loc1199)
      %31900 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31901 = stablehlo.broadcast_in_dim %31900, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1198)
      %31902 = stablehlo.add %31901, %31899 : tensor<1x1x2048x8192xbf16> loc(#loc1198)
      %31903 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %31904 = stablehlo.broadcast_in_dim %31903, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31905 = stablehlo.multiply %31904, %31902 : tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31906 = stablehlo.multiply %31889, %31905 : tensor<1x1x2048x8192xbf16> loc(#loc1197)
      %31907 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1200)
      %31908 = stablehlo.convert %31907 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1201)
      %31909 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31910 = stablehlo.broadcast_in_dim %31909, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1202)
      %31911 = stablehlo.subtract %31910, %31908 : tensor<1x1x2048x1xf32> loc(#loc1202)
      %31912 = stablehlo.convert %31911 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1203)
      %31913 = stablehlo.broadcast_in_dim %31912, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1204)
      %31914 = stablehlo.multiply %31906, %31913 : tensor<1x1x2048x8192xbf16> loc(#loc1204)
      %31915 = func.call @_where_303(%30042, %iterArg_292, %iterArg_293) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %31916 = func.call @_einsum_304(%31914, %31915) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1205)
      %31917 = arith.select %ptrue, %31916, %31916 : tensor<1x1x2048x2048xbf16>
      %31918 = func.call @_where_305(%30042, %iterArg_294, %iterArg_295) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31919 = stablehlo.broadcast_in_dim %31918, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1207)
      %31920 = stablehlo.broadcast_in_dim %31919, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1208)
      %31921 = stablehlo.add %31917, %31920 : tensor<1x1x2048x2048xbf16> loc(#loc1208)
      %31922 = stablehlo.convert %31907 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1201)
      %31923 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %31924 = stablehlo.broadcast_in_dim %31923, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1202)
      %31925 = stablehlo.subtract %31924, %31922 : tensor<1x1x2048x1xf32> loc(#loc1202)
      %31926 = stablehlo.convert %31925 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1209)
      %31927 = stablehlo.broadcast_in_dim %31926, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1210)
      %31928 = stablehlo.multiply %31921, %31927 : tensor<1x1x2048x2048xbf16> loc(#loc1210)
      %31929 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31930 = stablehlo.broadcast_in_dim %31929, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1211)
      %31931 = stablehlo.multiply %31928, %31930 : tensor<1x1x2048x2048xbf16> loc(#loc1211)
      %31932 = stablehlo.add %31844, %31931 : tensor<1x1x2048x2048xbf16> loc(#loc1212)
      %31933 = stablehlo.convert %31932 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1213)
      %31934 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1214)
      %31935 = stablehlo.reduce(%31933 init: %31934) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1214)
        stablehlo.return %33956 : tensor<f32> loc(#loc1214)
      } loc(#loc1214)
      %31936 = stablehlo.broadcast_in_dim %31935, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1215)
      %31937 = stablehlo.convert %31936 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1216)
      %31938 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31939 = stablehlo.broadcast_in_dim %31938, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1217)
      %31940 = stablehlo.divide %31937, %31939 : tensor<1x1x2048x1xbf16> loc(#loc1217)
      %31941 = stablehlo.broadcast_in_dim %31940, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1218)
      %31942 = stablehlo.subtract %31932, %31941 : tensor<1x1x2048x2048xbf16> loc(#loc1218)
      %31943 = stablehlo.broadcast_in_dim %31940, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1218)
      %31944 = stablehlo.subtract %31932, %31943 : tensor<1x1x2048x2048xbf16> loc(#loc1218)
      %31945 = stablehlo.multiply %31944, %31944 : tensor<1x1x2048x2048xbf16> loc(#loc1219)
      %31946 = stablehlo.convert %31945 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1213)
      %31947 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1214)
      %31948 = stablehlo.reduce(%31946 init: %31947) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1214)
        stablehlo.return %33956 : tensor<f32> loc(#loc1214)
      } loc(#loc1214)
      %31949 = stablehlo.broadcast_in_dim %31948, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1215)
      %31950 = stablehlo.convert %31949 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1216)
      %31951 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %31952 = stablehlo.broadcast_in_dim %31951, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1217)
      %31953 = stablehlo.divide %31950, %31952 : tensor<1x1x2048x1xbf16> loc(#loc1217)
      %31954 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %31955 = stablehlo.broadcast_in_dim %31954, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1220)
      %31956 = stablehlo.add %31953, %31955 : tensor<1x1x2048x1xbf16> loc(#loc1220)
      %31957 = stablehlo.rsqrt %31956 : tensor<1x1x2048x1xbf16> loc(#loc1221)
      %31958 = stablehlo.broadcast_in_dim %31957, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1222)
      %31959 = stablehlo.multiply %31942, %31958 : tensor<1x1x2048x2048xbf16> loc(#loc1222)
      %31960 = func.call @_where_306(%30042, %iterArg_296, %iterArg_297) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31961 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %31962 = stablehlo.broadcast_in_dim %31961, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1223)
      %31963 = stablehlo.add %31962, %31960 : tensor<1x2048xbf16> loc(#loc1223)
      %31964 = stablehlo.broadcast_in_dim %31963, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1224)
      %31965 = stablehlo.broadcast_in_dim %31964, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1225)
      %31966 = stablehlo.multiply %31959, %31965 : tensor<1x1x2048x2048xbf16> loc(#loc1225)
      %31967 = func.call @_where_307(%30042, %iterArg_298, %iterArg_299) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %31968 = stablehlo.broadcast_in_dim %31967, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1226)
      %31969 = stablehlo.broadcast_in_dim %31968, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1227)
      %31970 = stablehlo.add %31966, %31969 : tensor<1x1x2048x2048xbf16> loc(#loc1227)
      %31971 = func.call @_where_308(%30042, %iterArg_300, %iterArg_301) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %31972 = func.call @_einsum_309(%31970, %31971) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1228)
      %31973 = "stablehlo.slice"(%31972) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1229)
      %31974 = stablehlo.reshape %31973 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1230)
      %31975 = arith.select %ptrue, %31974, %31974 : tensor<1x1x2048x32x64xbf16>
      %31976 = func.call @_where_310(%30042, %iterArg_302, %iterArg_303) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %31977 = func.call @softplus_311(%31976, %iterArg_304) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1232)
      %31978 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %31979 = stablehlo.broadcast_in_dim %31978, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1233)
      %31980 = stablehlo.multiply %31979, %31977 : tensor<1x64xbf16> loc(#loc1233)
      %31981 = stablehlo.broadcast_in_dim %31980, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1234)
      %31982 = stablehlo.broadcast_in_dim %31981, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1235)
      %31983 = stablehlo.multiply %31975, %31982 : tensor<1x1x2048x32x64xbf16> loc(#loc1235)
      %31984 = "stablehlo.slice"(%31972) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1236)
      %31985 = stablehlo.reshape %31984 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1230)
      %31986 = arith.select %ptrue, %31985, %31985 : tensor<1x1x2048x32x64xbf16>
      %31987 = func.call @_einsum_313(%31983, %31986) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1237)
      %31988 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31989 = stablehlo.broadcast_in_dim %31988, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1238)
      %31990 = stablehlo.divide %31987, %31989 : tensor<1x1x32x2048x2048xbf16> loc(#loc1238)
      %31991 = stablehlo.tanh %31990 : tensor<1x1x32x2048x2048xbf16> loc(#loc1239)
      %31992 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %31993 = stablehlo.broadcast_in_dim %31992, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1240)
      %31994 = stablehlo.multiply %31993, %31991 : tensor<1x1x32x2048x2048xbf16> loc(#loc1240)
      %31995 = stablehlo.convert %31994 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1241)
      %31996 = func.call @_where_314(%29964, %31995, %iterArg_305) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1242)
      %31997 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1243)
      %31998 = stablehlo.reduce(%31996 init: %31997) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1243)
        stablehlo.return %33956 : tensor<f32> loc(#loc1243)
      } loc(#loc1243)
      %31999 = stablehlo.broadcast_in_dim %31998, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1244)
      %32000 = stablehlo.broadcast_in_dim %31999, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1245)
      %32001 = stablehlo.subtract %31996, %32000 : tensor<1x1x32x2048x2048xf32> loc(#loc1245)
      %32002 = stablehlo.exponential %32001 : tensor<1x1x32x2048x2048xf32> loc(#loc1246)
      %32003 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1247)
      %32004 = stablehlo.reduce(%32002 init: %32003) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1247)
        stablehlo.return %33956 : tensor<f32> loc(#loc1247)
      } loc(#loc1247)
      %32005 = stablehlo.broadcast_in_dim %32004, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1244)
      %32006 = stablehlo.broadcast_in_dim %32005, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1248)
      %32007 = stablehlo.divide %32002, %32006 : tensor<1x1x32x2048x2048xf32> loc(#loc1248)
      %32008 = stablehlo.convert %32007 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1249)
      %32009 = "stablehlo.slice"(%31972) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1250)
      %32010 = stablehlo.reshape %32009 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1230)
      %32011 = arith.select %ptrue, %32010, %32010 : tensor<1x1x2048x32x64xbf16>
      %32012 = func.call @_einsum_315(%32008, %32011) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1251)
      %32013 = arith.select %ptrue, %32012, %32012 : tensor<1x1x2048x32x64xbf16>
      %32014 = func.call @_where_316(%30042, %iterArg_306, %iterArg_307) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %32015 = func.call @_einsum_317(%32013, %32014) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1252)
      %32016 = arith.select %ptrue, %32015, %32015 : tensor<1x1x2048x2048xbf16>
      %32017 = stablehlo.add %32016, %31932 : tensor<1x1x2048x2048xbf16> loc(#loc1254)
      %32018 = stablehlo.convert %32017 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1255)
      %32019 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1256)
      %32020 = stablehlo.reduce(%32018 init: %32019) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1256)
        stablehlo.return %33956 : tensor<f32> loc(#loc1256)
      } loc(#loc1256)
      %32021 = stablehlo.broadcast_in_dim %32020, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1257)
      %32022 = stablehlo.convert %32021 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1258)
      %32023 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32024 = stablehlo.broadcast_in_dim %32023, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1259)
      %32025 = stablehlo.divide %32022, %32024 : tensor<1x1x2048x1xbf16> loc(#loc1259)
      %32026 = stablehlo.broadcast_in_dim %32025, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1260)
      %32027 = stablehlo.subtract %32017, %32026 : tensor<1x1x2048x2048xbf16> loc(#loc1260)
      %32028 = stablehlo.broadcast_in_dim %32025, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1260)
      %32029 = stablehlo.subtract %32017, %32028 : tensor<1x1x2048x2048xbf16> loc(#loc1260)
      %32030 = stablehlo.multiply %32029, %32029 : tensor<1x1x2048x2048xbf16> loc(#loc1261)
      %32031 = stablehlo.convert %32030 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1255)
      %32032 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1256)
      %32033 = stablehlo.reduce(%32031 init: %32032) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1256)
        stablehlo.return %33956 : tensor<f32> loc(#loc1256)
      } loc(#loc1256)
      %32034 = stablehlo.broadcast_in_dim %32033, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1257)
      %32035 = stablehlo.convert %32034 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1258)
      %32036 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32037 = stablehlo.broadcast_in_dim %32036, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1259)
      %32038 = stablehlo.divide %32035, %32037 : tensor<1x1x2048x1xbf16> loc(#loc1259)
      %32039 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32040 = stablehlo.broadcast_in_dim %32039, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1262)
      %32041 = stablehlo.add %32038, %32040 : tensor<1x1x2048x1xbf16> loc(#loc1262)
      %32042 = stablehlo.rsqrt %32041 : tensor<1x1x2048x1xbf16> loc(#loc1263)
      %32043 = stablehlo.broadcast_in_dim %32042, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1264)
      %32044 = stablehlo.multiply %32027, %32043 : tensor<1x1x2048x2048xbf16> loc(#loc1264)
      %32045 = func.call @_where_318(%30042, %iterArg_308, %iterArg_309) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32046 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32047 = stablehlo.broadcast_in_dim %32046, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1265)
      %32048 = stablehlo.add %32047, %32045 : tensor<1x2048xbf16> loc(#loc1265)
      %32049 = stablehlo.broadcast_in_dim %32048, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1266)
      %32050 = stablehlo.broadcast_in_dim %32049, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1267)
      %32051 = stablehlo.multiply %32044, %32050 : tensor<1x1x2048x2048xbf16> loc(#loc1267)
      %32052 = func.call @_where_319(%30042, %iterArg_310, %iterArg_311) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32053 = stablehlo.broadcast_in_dim %32052, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1268)
      %32054 = stablehlo.broadcast_in_dim %32053, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1269)
      %32055 = stablehlo.add %32051, %32054 : tensor<1x1x2048x2048xbf16> loc(#loc1269)
      %32056 = func.call @_where_320(%30042, %iterArg_312, %iterArg_313) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %32057 = func.call @_einsum_321(%32055, %32056) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1270)
      %32058 = arith.select %ptrue, %32057, %32057 : tensor<1x1x2048x8192xbf16>
      %32059 = func.call @_where_322(%30042, %iterArg_314, %iterArg_315) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %32060 = stablehlo.broadcast_in_dim %32059, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1272)
      %32061 = stablehlo.broadcast_in_dim %32060, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1273)
      %32062 = stablehlo.add %32058, %32061 : tensor<1x1x2048x8192xbf16> loc(#loc1273)
      %32063 = stablehlo.multiply %32062, %32062 : tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32064 = stablehlo.multiply %32062, %32063 : tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32065 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %32066 = stablehlo.broadcast_in_dim %32065, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32067 = stablehlo.multiply %32066, %32064 : tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32068 = stablehlo.add %32062, %32067 : tensor<1x1x2048x8192xbf16> loc(#loc1275)
      %32069 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %32070 = stablehlo.broadcast_in_dim %32069, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32071 = stablehlo.multiply %32070, %32068 : tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32072 = stablehlo.tanh %32071 : tensor<1x1x2048x8192xbf16> loc(#loc1276)
      %32073 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32074 = stablehlo.broadcast_in_dim %32073, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1275)
      %32075 = stablehlo.add %32074, %32072 : tensor<1x1x2048x8192xbf16> loc(#loc1275)
      %32076 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %32077 = stablehlo.broadcast_in_dim %32076, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32078 = stablehlo.multiply %32077, %32075 : tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32079 = stablehlo.multiply %32062, %32078 : tensor<1x1x2048x8192xbf16> loc(#loc1274)
      %32080 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1277)
      %32081 = stablehlo.convert %32080 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1278)
      %32082 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32083 = stablehlo.broadcast_in_dim %32082, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1279)
      %32084 = stablehlo.subtract %32083, %32081 : tensor<1x1x2048x1xf32> loc(#loc1279)
      %32085 = stablehlo.convert %32084 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1280)
      %32086 = stablehlo.broadcast_in_dim %32085, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1281)
      %32087 = stablehlo.multiply %32079, %32086 : tensor<1x1x2048x8192xbf16> loc(#loc1281)
      %32088 = func.call @_where_323(%30042, %iterArg_316, %iterArg_317) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %32089 = func.call @_einsum_324(%32087, %32088) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1282)
      %32090 = arith.select %ptrue, %32089, %32089 : tensor<1x1x2048x2048xbf16>
      %32091 = func.call @_where_325(%30042, %iterArg_318, %iterArg_319) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32092 = stablehlo.broadcast_in_dim %32091, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1284)
      %32093 = stablehlo.broadcast_in_dim %32092, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1285)
      %32094 = stablehlo.add %32090, %32093 : tensor<1x1x2048x2048xbf16> loc(#loc1285)
      %32095 = stablehlo.convert %32080 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1278)
      %32096 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32097 = stablehlo.broadcast_in_dim %32096, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1279)
      %32098 = stablehlo.subtract %32097, %32095 : tensor<1x1x2048x1xf32> loc(#loc1279)
      %32099 = stablehlo.convert %32098 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1286)
      %32100 = stablehlo.broadcast_in_dim %32099, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1287)
      %32101 = stablehlo.multiply %32094, %32100 : tensor<1x1x2048x2048xbf16> loc(#loc1287)
      %32102 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32103 = stablehlo.broadcast_in_dim %32102, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1288)
      %32104 = stablehlo.multiply %32101, %32103 : tensor<1x1x2048x2048xbf16> loc(#loc1288)
      %32105 = stablehlo.add %32017, %32104 : tensor<1x1x2048x2048xbf16> loc(#loc1289)
      %32106 = stablehlo.convert %32105 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1290)
      %32107 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1291)
      %32108 = stablehlo.reduce(%32106 init: %32107) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1291)
        stablehlo.return %33956 : tensor<f32> loc(#loc1291)
      } loc(#loc1291)
      %32109 = stablehlo.broadcast_in_dim %32108, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1292)
      %32110 = stablehlo.convert %32109 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1293)
      %32111 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32112 = stablehlo.broadcast_in_dim %32111, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1294)
      %32113 = stablehlo.divide %32110, %32112 : tensor<1x1x2048x1xbf16> loc(#loc1294)
      %32114 = stablehlo.broadcast_in_dim %32113, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1295)
      %32115 = stablehlo.subtract %32105, %32114 : tensor<1x1x2048x2048xbf16> loc(#loc1295)
      %32116 = stablehlo.broadcast_in_dim %32113, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1295)
      %32117 = stablehlo.subtract %32105, %32116 : tensor<1x1x2048x2048xbf16> loc(#loc1295)
      %32118 = stablehlo.multiply %32117, %32117 : tensor<1x1x2048x2048xbf16> loc(#loc1296)
      %32119 = stablehlo.convert %32118 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1290)
      %32120 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1291)
      %32121 = stablehlo.reduce(%32119 init: %32120) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1291)
        stablehlo.return %33956 : tensor<f32> loc(#loc1291)
      } loc(#loc1291)
      %32122 = stablehlo.broadcast_in_dim %32121, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1292)
      %32123 = stablehlo.convert %32122 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1293)
      %32124 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32125 = stablehlo.broadcast_in_dim %32124, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1294)
      %32126 = stablehlo.divide %32123, %32125 : tensor<1x1x2048x1xbf16> loc(#loc1294)
      %32127 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32128 = stablehlo.broadcast_in_dim %32127, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1297)
      %32129 = stablehlo.add %32126, %32128 : tensor<1x1x2048x1xbf16> loc(#loc1297)
      %32130 = stablehlo.rsqrt %32129 : tensor<1x1x2048x1xbf16> loc(#loc1298)
      %32131 = stablehlo.broadcast_in_dim %32130, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1299)
      %32132 = stablehlo.multiply %32115, %32131 : tensor<1x1x2048x2048xbf16> loc(#loc1299)
      %32133 = func.call @_where_326(%30042, %iterArg_320, %iterArg_321) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32134 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32135 = stablehlo.broadcast_in_dim %32134, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1300)
      %32136 = stablehlo.add %32135, %32133 : tensor<1x2048xbf16> loc(#loc1300)
      %32137 = stablehlo.broadcast_in_dim %32136, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1301)
      %32138 = stablehlo.broadcast_in_dim %32137, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1302)
      %32139 = stablehlo.multiply %32132, %32138 : tensor<1x1x2048x2048xbf16> loc(#loc1302)
      %32140 = func.call @_where_327(%30042, %iterArg_322, %iterArg_323) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32141 = stablehlo.broadcast_in_dim %32140, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1303)
      %32142 = stablehlo.broadcast_in_dim %32141, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1304)
      %32143 = stablehlo.add %32139, %32142 : tensor<1x1x2048x2048xbf16> loc(#loc1304)
      %32144 = func.call @_where_328(%30042, %iterArg_324, %iterArg_325) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %32145 = func.call @_einsum_329(%32143, %32144) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1305)
      %32146 = "stablehlo.slice"(%32145) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1306)
      %32147 = stablehlo.reshape %32146 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1307)
      %32148 = arith.select %ptrue, %32147, %32147 : tensor<1x1x2048x32x64xbf16>
      %32149 = func.call @_where_330(%30042, %iterArg_326, %iterArg_327) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %32150 = func.call @softplus_331(%32149, %iterArg_328) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1309)
      %32151 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %32152 = stablehlo.broadcast_in_dim %32151, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1310)
      %32153 = stablehlo.multiply %32152, %32150 : tensor<1x64xbf16> loc(#loc1310)
      %32154 = stablehlo.broadcast_in_dim %32153, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1311)
      %32155 = stablehlo.broadcast_in_dim %32154, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1312)
      %32156 = stablehlo.multiply %32148, %32155 : tensor<1x1x2048x32x64xbf16> loc(#loc1312)
      %32157 = "stablehlo.slice"(%32145) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1313)
      %32158 = stablehlo.reshape %32157 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1307)
      %32159 = arith.select %ptrue, %32158, %32158 : tensor<1x1x2048x32x64xbf16>
      %32160 = func.call @_einsum_333(%32156, %32159) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1314)
      %32161 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32162 = stablehlo.broadcast_in_dim %32161, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1315)
      %32163 = stablehlo.divide %32160, %32162 : tensor<1x1x32x2048x2048xbf16> loc(#loc1315)
      %32164 = stablehlo.tanh %32163 : tensor<1x1x32x2048x2048xbf16> loc(#loc1316)
      %32165 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32166 = stablehlo.broadcast_in_dim %32165, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1317)
      %32167 = stablehlo.multiply %32166, %32164 : tensor<1x1x32x2048x2048xbf16> loc(#loc1317)
      %32168 = stablehlo.convert %32167 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1318)
      %32169 = func.call @_where_334(%29961, %32168, %iterArg_329) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1319)
      %32170 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1320)
      %32171 = stablehlo.reduce(%32169 init: %32170) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1320)
        stablehlo.return %33956 : tensor<f32> loc(#loc1320)
      } loc(#loc1320)
      %32172 = stablehlo.broadcast_in_dim %32171, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1321)
      %32173 = stablehlo.broadcast_in_dim %32172, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1322)
      %32174 = stablehlo.subtract %32169, %32173 : tensor<1x1x32x2048x2048xf32> loc(#loc1322)
      %32175 = stablehlo.exponential %32174 : tensor<1x1x32x2048x2048xf32> loc(#loc1323)
      %32176 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1324)
      %32177 = stablehlo.reduce(%32175 init: %32176) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1324)
        stablehlo.return %33956 : tensor<f32> loc(#loc1324)
      } loc(#loc1324)
      %32178 = stablehlo.broadcast_in_dim %32177, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1321)
      %32179 = stablehlo.broadcast_in_dim %32178, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1325)
      %32180 = stablehlo.divide %32175, %32179 : tensor<1x1x32x2048x2048xf32> loc(#loc1325)
      %32181 = stablehlo.convert %32180 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1326)
      %32182 = "stablehlo.slice"(%32145) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1327)
      %32183 = stablehlo.reshape %32182 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1307)
      %32184 = arith.select %ptrue, %32183, %32183 : tensor<1x1x2048x32x64xbf16>
      %32185 = func.call @_einsum_335(%32181, %32184) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1328)
      %32186 = arith.select %ptrue, %32185, %32185 : tensor<1x1x2048x32x64xbf16>
      %32187 = func.call @_where_336(%30042, %iterArg_330, %iterArg_331) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %32188 = func.call @_einsum_337(%32186, %32187) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1329)
      %32189 = arith.select %ptrue, %32188, %32188 : tensor<1x1x2048x2048xbf16>
      %32190 = stablehlo.add %32189, %32105 : tensor<1x1x2048x2048xbf16> loc(#loc1331)
      %32191 = stablehlo.convert %32190 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1332)
      %32192 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1333)
      %32193 = stablehlo.reduce(%32191 init: %32192) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1333)
        stablehlo.return %33956 : tensor<f32> loc(#loc1333)
      } loc(#loc1333)
      %32194 = stablehlo.broadcast_in_dim %32193, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1334)
      %32195 = stablehlo.convert %32194 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1335)
      %32196 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32197 = stablehlo.broadcast_in_dim %32196, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1336)
      %32198 = stablehlo.divide %32195, %32197 : tensor<1x1x2048x1xbf16> loc(#loc1336)
      %32199 = stablehlo.broadcast_in_dim %32198, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1337)
      %32200 = stablehlo.subtract %32190, %32199 : tensor<1x1x2048x2048xbf16> loc(#loc1337)
      %32201 = stablehlo.broadcast_in_dim %32198, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1337)
      %32202 = stablehlo.subtract %32190, %32201 : tensor<1x1x2048x2048xbf16> loc(#loc1337)
      %32203 = stablehlo.multiply %32202, %32202 : tensor<1x1x2048x2048xbf16> loc(#loc1338)
      %32204 = stablehlo.convert %32203 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1332)
      %32205 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1333)
      %32206 = stablehlo.reduce(%32204 init: %32205) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1333)
        stablehlo.return %33956 : tensor<f32> loc(#loc1333)
      } loc(#loc1333)
      %32207 = stablehlo.broadcast_in_dim %32206, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1334)
      %32208 = stablehlo.convert %32207 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1335)
      %32209 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32210 = stablehlo.broadcast_in_dim %32209, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1336)
      %32211 = stablehlo.divide %32208, %32210 : tensor<1x1x2048x1xbf16> loc(#loc1336)
      %32212 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32213 = stablehlo.broadcast_in_dim %32212, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1339)
      %32214 = stablehlo.add %32211, %32213 : tensor<1x1x2048x1xbf16> loc(#loc1339)
      %32215 = stablehlo.rsqrt %32214 : tensor<1x1x2048x1xbf16> loc(#loc1340)
      %32216 = stablehlo.broadcast_in_dim %32215, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1341)
      %32217 = stablehlo.multiply %32200, %32216 : tensor<1x1x2048x2048xbf16> loc(#loc1341)
      %32218 = func.call @_where_338(%30042, %iterArg_332, %iterArg_333) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32219 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32220 = stablehlo.broadcast_in_dim %32219, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1342)
      %32221 = stablehlo.add %32220, %32218 : tensor<1x2048xbf16> loc(#loc1342)
      %32222 = stablehlo.broadcast_in_dim %32221, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1343)
      %32223 = stablehlo.broadcast_in_dim %32222, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1344)
      %32224 = stablehlo.multiply %32217, %32223 : tensor<1x1x2048x2048xbf16> loc(#loc1344)
      %32225 = func.call @_where_339(%30042, %iterArg_334, %iterArg_335) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32226 = stablehlo.broadcast_in_dim %32225, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1345)
      %32227 = stablehlo.broadcast_in_dim %32226, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1346)
      %32228 = stablehlo.add %32224, %32227 : tensor<1x1x2048x2048xbf16> loc(#loc1346)
      %32229 = func.call @_where_340(%30042, %iterArg_336, %iterArg_337) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %32230 = func.call @_einsum_341(%32228, %32229) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1347)
      %32231 = arith.select %ptrue, %32230, %32230 : tensor<1x1x2048x8192xbf16>
      %32232 = func.call @_where_342(%30042, %iterArg_338, %iterArg_339) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %32233 = stablehlo.broadcast_in_dim %32232, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1349)
      %32234 = stablehlo.broadcast_in_dim %32233, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1350)
      %32235 = stablehlo.add %32231, %32234 : tensor<1x1x2048x8192xbf16> loc(#loc1350)
      %32236 = stablehlo.multiply %32235, %32235 : tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32237 = stablehlo.multiply %32235, %32236 : tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32238 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %32239 = stablehlo.broadcast_in_dim %32238, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32240 = stablehlo.multiply %32239, %32237 : tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32241 = stablehlo.add %32235, %32240 : tensor<1x1x2048x8192xbf16> loc(#loc1352)
      %32242 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %32243 = stablehlo.broadcast_in_dim %32242, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32244 = stablehlo.multiply %32243, %32241 : tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32245 = stablehlo.tanh %32244 : tensor<1x1x2048x8192xbf16> loc(#loc1353)
      %32246 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32247 = stablehlo.broadcast_in_dim %32246, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1352)
      %32248 = stablehlo.add %32247, %32245 : tensor<1x1x2048x8192xbf16> loc(#loc1352)
      %32249 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %32250 = stablehlo.broadcast_in_dim %32249, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32251 = stablehlo.multiply %32250, %32248 : tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32252 = stablehlo.multiply %32235, %32251 : tensor<1x1x2048x8192xbf16> loc(#loc1351)
      %32253 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1354)
      %32254 = stablehlo.convert %32253 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1355)
      %32255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32256 = stablehlo.broadcast_in_dim %32255, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1356)
      %32257 = stablehlo.subtract %32256, %32254 : tensor<1x1x2048x1xf32> loc(#loc1356)
      %32258 = stablehlo.convert %32257 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1357)
      %32259 = stablehlo.broadcast_in_dim %32258, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1358)
      %32260 = stablehlo.multiply %32252, %32259 : tensor<1x1x2048x8192xbf16> loc(#loc1358)
      %32261 = func.call @_where_343(%30042, %iterArg_340, %iterArg_341) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %32262 = func.call @_einsum_344(%32260, %32261) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1359)
      %32263 = arith.select %ptrue, %32262, %32262 : tensor<1x1x2048x2048xbf16>
      %32264 = func.call @_where_345(%30042, %iterArg_342, %iterArg_343) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32265 = stablehlo.broadcast_in_dim %32264, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1361)
      %32266 = stablehlo.broadcast_in_dim %32265, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1362)
      %32267 = stablehlo.add %32263, %32266 : tensor<1x1x2048x2048xbf16> loc(#loc1362)
      %32268 = stablehlo.convert %32253 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1355)
      %32269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32270 = stablehlo.broadcast_in_dim %32269, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1356)
      %32271 = stablehlo.subtract %32270, %32268 : tensor<1x1x2048x1xf32> loc(#loc1356)
      %32272 = stablehlo.convert %32271 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1363)
      %32273 = stablehlo.broadcast_in_dim %32272, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1364)
      %32274 = stablehlo.multiply %32267, %32273 : tensor<1x1x2048x2048xbf16> loc(#loc1364)
      %32275 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32276 = stablehlo.broadcast_in_dim %32275, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1365)
      %32277 = stablehlo.multiply %32274, %32276 : tensor<1x1x2048x2048xbf16> loc(#loc1365)
      %32278 = stablehlo.add %32190, %32277 : tensor<1x1x2048x2048xbf16> loc(#loc1366)
      %32279 = stablehlo.convert %32278 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1367)
      %32280 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1368)
      %32281 = stablehlo.reduce(%32279 init: %32280) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1368)
        stablehlo.return %33956 : tensor<f32> loc(#loc1368)
      } loc(#loc1368)
      %32282 = stablehlo.broadcast_in_dim %32281, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1369)
      %32283 = stablehlo.convert %32282 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1370)
      %32284 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32285 = stablehlo.broadcast_in_dim %32284, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1371)
      %32286 = stablehlo.divide %32283, %32285 : tensor<1x1x2048x1xbf16> loc(#loc1371)
      %32287 = stablehlo.broadcast_in_dim %32286, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1372)
      %32288 = stablehlo.subtract %32278, %32287 : tensor<1x1x2048x2048xbf16> loc(#loc1372)
      %32289 = stablehlo.broadcast_in_dim %32286, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1372)
      %32290 = stablehlo.subtract %32278, %32289 : tensor<1x1x2048x2048xbf16> loc(#loc1372)
      %32291 = stablehlo.multiply %32290, %32290 : tensor<1x1x2048x2048xbf16> loc(#loc1373)
      %32292 = stablehlo.convert %32291 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1367)
      %32293 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1368)
      %32294 = stablehlo.reduce(%32292 init: %32293) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1368)
        stablehlo.return %33956 : tensor<f32> loc(#loc1368)
      } loc(#loc1368)
      %32295 = stablehlo.broadcast_in_dim %32294, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1369)
      %32296 = stablehlo.convert %32295 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1370)
      %32297 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32298 = stablehlo.broadcast_in_dim %32297, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1371)
      %32299 = stablehlo.divide %32296, %32298 : tensor<1x1x2048x1xbf16> loc(#loc1371)
      %32300 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32301 = stablehlo.broadcast_in_dim %32300, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1374)
      %32302 = stablehlo.add %32299, %32301 : tensor<1x1x2048x1xbf16> loc(#loc1374)
      %32303 = stablehlo.rsqrt %32302 : tensor<1x1x2048x1xbf16> loc(#loc1375)
      %32304 = stablehlo.broadcast_in_dim %32303, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1376)
      %32305 = stablehlo.multiply %32288, %32304 : tensor<1x1x2048x2048xbf16> loc(#loc1376)
      %32306 = func.call @_where_346(%30042, %iterArg_344, %iterArg_345) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32307 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32308 = stablehlo.broadcast_in_dim %32307, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1377)
      %32309 = stablehlo.add %32308, %32306 : tensor<1x2048xbf16> loc(#loc1377)
      %32310 = stablehlo.broadcast_in_dim %32309, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1378)
      %32311 = stablehlo.broadcast_in_dim %32310, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1379)
      %32312 = stablehlo.multiply %32305, %32311 : tensor<1x1x2048x2048xbf16> loc(#loc1379)
      %32313 = func.call @_where_347(%30042, %iterArg_346, %iterArg_347) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32314 = stablehlo.broadcast_in_dim %32313, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1380)
      %32315 = stablehlo.broadcast_in_dim %32314, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1381)
      %32316 = stablehlo.add %32312, %32315 : tensor<1x1x2048x2048xbf16> loc(#loc1381)
      %32317 = func.call @_where_348(%30042, %iterArg_348, %iterArg_349) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %32318 = func.call @_einsum_349(%32316, %32317) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1382)
      %32319 = "stablehlo.slice"(%32318) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1383)
      %32320 = stablehlo.reshape %32319 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1384)
      %32321 = arith.select %ptrue, %32320, %32320 : tensor<1x1x2048x32x64xbf16>
      %32322 = func.call @_where_350(%30042, %iterArg_350, %iterArg_351) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %32323 = func.call @softplus_351(%32322, %iterArg_352) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1386)
      %32324 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %32325 = stablehlo.broadcast_in_dim %32324, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1387)
      %32326 = stablehlo.multiply %32325, %32323 : tensor<1x64xbf16> loc(#loc1387)
      %32327 = stablehlo.broadcast_in_dim %32326, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1388)
      %32328 = stablehlo.broadcast_in_dim %32327, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1389)
      %32329 = stablehlo.multiply %32321, %32328 : tensor<1x1x2048x32x64xbf16> loc(#loc1389)
      %32330 = "stablehlo.slice"(%32318) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1390)
      %32331 = stablehlo.reshape %32330 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1384)
      %32332 = arith.select %ptrue, %32331, %32331 : tensor<1x1x2048x32x64xbf16>
      %32333 = func.call @_einsum_353(%32329, %32332) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1391)
      %32334 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32335 = stablehlo.broadcast_in_dim %32334, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1392)
      %32336 = stablehlo.divide %32333, %32335 : tensor<1x1x32x2048x2048xbf16> loc(#loc1392)
      %32337 = stablehlo.tanh %32336 : tensor<1x1x32x2048x2048xbf16> loc(#loc1393)
      %32338 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32339 = stablehlo.broadcast_in_dim %32338, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1394)
      %32340 = stablehlo.multiply %32339, %32337 : tensor<1x1x32x2048x2048xbf16> loc(#loc1394)
      %32341 = stablehlo.convert %32340 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1395)
      %32342 = func.call @_where_354(%29958, %32341, %iterArg_353) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1396)
      %32343 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1397)
      %32344 = stablehlo.reduce(%32342 init: %32343) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1397)
        stablehlo.return %33956 : tensor<f32> loc(#loc1397)
      } loc(#loc1397)
      %32345 = stablehlo.broadcast_in_dim %32344, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1398)
      %32346 = stablehlo.broadcast_in_dim %32345, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1399)
      %32347 = stablehlo.subtract %32342, %32346 : tensor<1x1x32x2048x2048xf32> loc(#loc1399)
      %32348 = stablehlo.exponential %32347 : tensor<1x1x32x2048x2048xf32> loc(#loc1400)
      %32349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1401)
      %32350 = stablehlo.reduce(%32348 init: %32349) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1401)
        stablehlo.return %33956 : tensor<f32> loc(#loc1401)
      } loc(#loc1401)
      %32351 = stablehlo.broadcast_in_dim %32350, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1398)
      %32352 = stablehlo.broadcast_in_dim %32351, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1402)
      %32353 = stablehlo.divide %32348, %32352 : tensor<1x1x32x2048x2048xf32> loc(#loc1402)
      %32354 = stablehlo.convert %32353 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1403)
      %32355 = "stablehlo.slice"(%32318) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1404)
      %32356 = stablehlo.reshape %32355 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1384)
      %32357 = arith.select %ptrue, %32356, %32356 : tensor<1x1x2048x32x64xbf16>
      %32358 = func.call @_einsum_355(%32354, %32357) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1405)
      %32359 = arith.select %ptrue, %32358, %32358 : tensor<1x1x2048x32x64xbf16>
      %32360 = func.call @_where_356(%30042, %iterArg_354, %iterArg_355) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %32361 = func.call @_einsum_357(%32359, %32360) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1406)
      %32362 = arith.select %ptrue, %32361, %32361 : tensor<1x1x2048x2048xbf16>
      %32363 = stablehlo.add %32362, %32278 : tensor<1x1x2048x2048xbf16> loc(#loc1408)
      %32364 = stablehlo.convert %32363 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1409)
      %32365 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1410)
      %32366 = stablehlo.reduce(%32364 init: %32365) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1410)
        stablehlo.return %33956 : tensor<f32> loc(#loc1410)
      } loc(#loc1410)
      %32367 = stablehlo.broadcast_in_dim %32366, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1411)
      %32368 = stablehlo.convert %32367 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1412)
      %32369 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32370 = stablehlo.broadcast_in_dim %32369, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1413)
      %32371 = stablehlo.divide %32368, %32370 : tensor<1x1x2048x1xbf16> loc(#loc1413)
      %32372 = stablehlo.broadcast_in_dim %32371, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1414)
      %32373 = stablehlo.subtract %32363, %32372 : tensor<1x1x2048x2048xbf16> loc(#loc1414)
      %32374 = stablehlo.broadcast_in_dim %32371, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1414)
      %32375 = stablehlo.subtract %32363, %32374 : tensor<1x1x2048x2048xbf16> loc(#loc1414)
      %32376 = stablehlo.multiply %32375, %32375 : tensor<1x1x2048x2048xbf16> loc(#loc1415)
      %32377 = stablehlo.convert %32376 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1409)
      %32378 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1410)
      %32379 = stablehlo.reduce(%32377 init: %32378) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1410)
        stablehlo.return %33956 : tensor<f32> loc(#loc1410)
      } loc(#loc1410)
      %32380 = stablehlo.broadcast_in_dim %32379, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1411)
      %32381 = stablehlo.convert %32380 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1412)
      %32382 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32383 = stablehlo.broadcast_in_dim %32382, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1413)
      %32384 = stablehlo.divide %32381, %32383 : tensor<1x1x2048x1xbf16> loc(#loc1413)
      %32385 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32386 = stablehlo.broadcast_in_dim %32385, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1416)
      %32387 = stablehlo.add %32384, %32386 : tensor<1x1x2048x1xbf16> loc(#loc1416)
      %32388 = stablehlo.rsqrt %32387 : tensor<1x1x2048x1xbf16> loc(#loc1417)
      %32389 = stablehlo.broadcast_in_dim %32388, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1418)
      %32390 = stablehlo.multiply %32373, %32389 : tensor<1x1x2048x2048xbf16> loc(#loc1418)
      %32391 = func.call @_where_358(%30042, %iterArg_356, %iterArg_357) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32392 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32393 = stablehlo.broadcast_in_dim %32392, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1419)
      %32394 = stablehlo.add %32393, %32391 : tensor<1x2048xbf16> loc(#loc1419)
      %32395 = stablehlo.broadcast_in_dim %32394, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1420)
      %32396 = stablehlo.broadcast_in_dim %32395, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1421)
      %32397 = stablehlo.multiply %32390, %32396 : tensor<1x1x2048x2048xbf16> loc(#loc1421)
      %32398 = func.call @_where_359(%30042, %iterArg_358, %iterArg_359) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32399 = stablehlo.broadcast_in_dim %32398, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1422)
      %32400 = stablehlo.broadcast_in_dim %32399, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1423)
      %32401 = stablehlo.add %32397, %32400 : tensor<1x1x2048x2048xbf16> loc(#loc1423)
      %32402 = func.call @_where_360(%30042, %iterArg_360, %iterArg_361) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %32403 = func.call @_einsum_361(%32401, %32402) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1424)
      %32404 = arith.select %ptrue, %32403, %32403 : tensor<1x1x2048x8192xbf16>
      %32405 = func.call @_where_362(%30042, %iterArg_362, %iterArg_363) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %32406 = stablehlo.broadcast_in_dim %32405, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1426)
      %32407 = stablehlo.broadcast_in_dim %32406, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1427)
      %32408 = stablehlo.add %32404, %32407 : tensor<1x1x2048x8192xbf16> loc(#loc1427)
      %32409 = stablehlo.multiply %32408, %32408 : tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32410 = stablehlo.multiply %32408, %32409 : tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32411 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %32412 = stablehlo.broadcast_in_dim %32411, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32413 = stablehlo.multiply %32412, %32410 : tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32414 = stablehlo.add %32408, %32413 : tensor<1x1x2048x8192xbf16> loc(#loc1429)
      %32415 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %32416 = stablehlo.broadcast_in_dim %32415, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32417 = stablehlo.multiply %32416, %32414 : tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32418 = stablehlo.tanh %32417 : tensor<1x1x2048x8192xbf16> loc(#loc1430)
      %32419 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32420 = stablehlo.broadcast_in_dim %32419, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1429)
      %32421 = stablehlo.add %32420, %32418 : tensor<1x1x2048x8192xbf16> loc(#loc1429)
      %32422 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %32423 = stablehlo.broadcast_in_dim %32422, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32424 = stablehlo.multiply %32423, %32421 : tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32425 = stablehlo.multiply %32408, %32424 : tensor<1x1x2048x8192xbf16> loc(#loc1428)
      %32426 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1431)
      %32427 = stablehlo.convert %32426 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1432)
      %32428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32429 = stablehlo.broadcast_in_dim %32428, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1433)
      %32430 = stablehlo.subtract %32429, %32427 : tensor<1x1x2048x1xf32> loc(#loc1433)
      %32431 = stablehlo.convert %32430 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1434)
      %32432 = stablehlo.broadcast_in_dim %32431, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1435)
      %32433 = stablehlo.multiply %32425, %32432 : tensor<1x1x2048x8192xbf16> loc(#loc1435)
      %32434 = func.call @_where_363(%30042, %iterArg_364, %iterArg_365) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %32435 = func.call @_einsum_364(%32433, %32434) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1436)
      %32436 = arith.select %ptrue, %32435, %32435 : tensor<1x1x2048x2048xbf16>
      %32437 = func.call @_where_365(%30042, %iterArg_366, %iterArg_367) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32438 = stablehlo.broadcast_in_dim %32437, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1438)
      %32439 = stablehlo.broadcast_in_dim %32438, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1439)
      %32440 = stablehlo.add %32436, %32439 : tensor<1x1x2048x2048xbf16> loc(#loc1439)
      %32441 = stablehlo.convert %32426 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1432)
      %32442 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32443 = stablehlo.broadcast_in_dim %32442, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1433)
      %32444 = stablehlo.subtract %32443, %32441 : tensor<1x1x2048x1xf32> loc(#loc1433)
      %32445 = stablehlo.convert %32444 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1440)
      %32446 = stablehlo.broadcast_in_dim %32445, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1441)
      %32447 = stablehlo.multiply %32440, %32446 : tensor<1x1x2048x2048xbf16> loc(#loc1441)
      %32448 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32449 = stablehlo.broadcast_in_dim %32448, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1442)
      %32450 = stablehlo.multiply %32447, %32449 : tensor<1x1x2048x2048xbf16> loc(#loc1442)
      %32451 = stablehlo.add %32363, %32450 : tensor<1x1x2048x2048xbf16> loc(#loc1443)
      %32452 = stablehlo.convert %32451 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1444)
      %32453 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1445)
      %32454 = stablehlo.reduce(%32452 init: %32453) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1445)
        stablehlo.return %33956 : tensor<f32> loc(#loc1445)
      } loc(#loc1445)
      %32455 = stablehlo.broadcast_in_dim %32454, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1446)
      %32456 = stablehlo.convert %32455 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1447)
      %32457 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32458 = stablehlo.broadcast_in_dim %32457, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1448)
      %32459 = stablehlo.divide %32456, %32458 : tensor<1x1x2048x1xbf16> loc(#loc1448)
      %32460 = stablehlo.broadcast_in_dim %32459, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1449)
      %32461 = stablehlo.subtract %32451, %32460 : tensor<1x1x2048x2048xbf16> loc(#loc1449)
      %32462 = stablehlo.broadcast_in_dim %32459, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1449)
      %32463 = stablehlo.subtract %32451, %32462 : tensor<1x1x2048x2048xbf16> loc(#loc1449)
      %32464 = stablehlo.multiply %32463, %32463 : tensor<1x1x2048x2048xbf16> loc(#loc1450)
      %32465 = stablehlo.convert %32464 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1444)
      %32466 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1445)
      %32467 = stablehlo.reduce(%32465 init: %32466) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1445)
        stablehlo.return %33956 : tensor<f32> loc(#loc1445)
      } loc(#loc1445)
      %32468 = stablehlo.broadcast_in_dim %32467, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1446)
      %32469 = stablehlo.convert %32468 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1447)
      %32470 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32471 = stablehlo.broadcast_in_dim %32470, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1448)
      %32472 = stablehlo.divide %32469, %32471 : tensor<1x1x2048x1xbf16> loc(#loc1448)
      %32473 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32474 = stablehlo.broadcast_in_dim %32473, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1451)
      %32475 = stablehlo.add %32472, %32474 : tensor<1x1x2048x1xbf16> loc(#loc1451)
      %32476 = stablehlo.rsqrt %32475 : tensor<1x1x2048x1xbf16> loc(#loc1452)
      %32477 = stablehlo.broadcast_in_dim %32476, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1453)
      %32478 = stablehlo.multiply %32461, %32477 : tensor<1x1x2048x2048xbf16> loc(#loc1453)
      %32479 = func.call @_where_366(%30042, %iterArg_368, %iterArg_369) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32480 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32481 = stablehlo.broadcast_in_dim %32480, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1454)
      %32482 = stablehlo.add %32481, %32479 : tensor<1x2048xbf16> loc(#loc1454)
      %32483 = stablehlo.broadcast_in_dim %32482, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1455)
      %32484 = stablehlo.broadcast_in_dim %32483, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1456)
      %32485 = stablehlo.multiply %32478, %32484 : tensor<1x1x2048x2048xbf16> loc(#loc1456)
      %32486 = func.call @_where_367(%30042, %iterArg_370, %iterArg_371) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32487 = stablehlo.broadcast_in_dim %32486, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1457)
      %32488 = stablehlo.broadcast_in_dim %32487, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1458)
      %32489 = stablehlo.add %32485, %32488 : tensor<1x1x2048x2048xbf16> loc(#loc1458)
      %32490 = func.call @_where_368(%30042, %iterArg_372, %iterArg_373) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %32491 = func.call @_einsum_369(%32489, %32490) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1459)
      %32492 = "stablehlo.slice"(%32491) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1460)
      %32493 = stablehlo.reshape %32492 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1461)
      %32494 = arith.select %ptrue, %32493, %32493 : tensor<1x1x2048x32x64xbf16>
      %32495 = func.call @_where_370(%30042, %iterArg_374, %iterArg_375) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %32496 = func.call @softplus_371(%32495, %iterArg_376) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1463)
      %32497 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %32498 = stablehlo.broadcast_in_dim %32497, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1464)
      %32499 = stablehlo.multiply %32498, %32496 : tensor<1x64xbf16> loc(#loc1464)
      %32500 = stablehlo.broadcast_in_dim %32499, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1465)
      %32501 = stablehlo.broadcast_in_dim %32500, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1466)
      %32502 = stablehlo.multiply %32494, %32501 : tensor<1x1x2048x32x64xbf16> loc(#loc1466)
      %32503 = "stablehlo.slice"(%32491) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1467)
      %32504 = stablehlo.reshape %32503 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1461)
      %32505 = arith.select %ptrue, %32504, %32504 : tensor<1x1x2048x32x64xbf16>
      %32506 = func.call @_einsum_373(%32502, %32505) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1468)
      %32507 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32508 = stablehlo.broadcast_in_dim %32507, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1469)
      %32509 = stablehlo.divide %32506, %32508 : tensor<1x1x32x2048x2048xbf16> loc(#loc1469)
      %32510 = stablehlo.tanh %32509 : tensor<1x1x32x2048x2048xbf16> loc(#loc1470)
      %32511 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32512 = stablehlo.broadcast_in_dim %32511, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1471)
      %32513 = stablehlo.multiply %32512, %32510 : tensor<1x1x32x2048x2048xbf16> loc(#loc1471)
      %32514 = stablehlo.convert %32513 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1472)
      %32515 = func.call @_where_374(%29955, %32514, %iterArg_377) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1473)
      %32516 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1474)
      %32517 = stablehlo.reduce(%32515 init: %32516) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1474)
        stablehlo.return %33956 : tensor<f32> loc(#loc1474)
      } loc(#loc1474)
      %32518 = stablehlo.broadcast_in_dim %32517, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1475)
      %32519 = stablehlo.broadcast_in_dim %32518, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1476)
      %32520 = stablehlo.subtract %32515, %32519 : tensor<1x1x32x2048x2048xf32> loc(#loc1476)
      %32521 = stablehlo.exponential %32520 : tensor<1x1x32x2048x2048xf32> loc(#loc1477)
      %32522 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1478)
      %32523 = stablehlo.reduce(%32521 init: %32522) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1478)
        stablehlo.return %33956 : tensor<f32> loc(#loc1478)
      } loc(#loc1478)
      %32524 = stablehlo.broadcast_in_dim %32523, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1475)
      %32525 = stablehlo.broadcast_in_dim %32524, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1479)
      %32526 = stablehlo.divide %32521, %32525 : tensor<1x1x32x2048x2048xf32> loc(#loc1479)
      %32527 = stablehlo.convert %32526 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1480)
      %32528 = "stablehlo.slice"(%32491) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1481)
      %32529 = stablehlo.reshape %32528 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1461)
      %32530 = arith.select %ptrue, %32529, %32529 : tensor<1x1x2048x32x64xbf16>
      %32531 = func.call @_einsum_375(%32527, %32530) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1482)
      %32532 = arith.select %ptrue, %32531, %32531 : tensor<1x1x2048x32x64xbf16>
      %32533 = func.call @_where_376(%30042, %iterArg_378, %iterArg_379) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %32534 = func.call @_einsum_377(%32532, %32533) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1483)
      %32535 = arith.select %ptrue, %32534, %32534 : tensor<1x1x2048x2048xbf16>
      %32536 = stablehlo.add %32535, %32451 : tensor<1x1x2048x2048xbf16> loc(#loc1485)
      %32537 = stablehlo.convert %32536 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1486)
      %32538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1487)
      %32539 = stablehlo.reduce(%32537 init: %32538) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1487)
        stablehlo.return %33956 : tensor<f32> loc(#loc1487)
      } loc(#loc1487)
      %32540 = stablehlo.broadcast_in_dim %32539, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1488)
      %32541 = stablehlo.convert %32540 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1489)
      %32542 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32543 = stablehlo.broadcast_in_dim %32542, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1490)
      %32544 = stablehlo.divide %32541, %32543 : tensor<1x1x2048x1xbf16> loc(#loc1490)
      %32545 = stablehlo.broadcast_in_dim %32544, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1491)
      %32546 = stablehlo.subtract %32536, %32545 : tensor<1x1x2048x2048xbf16> loc(#loc1491)
      %32547 = stablehlo.broadcast_in_dim %32544, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1491)
      %32548 = stablehlo.subtract %32536, %32547 : tensor<1x1x2048x2048xbf16> loc(#loc1491)
      %32549 = stablehlo.multiply %32548, %32548 : tensor<1x1x2048x2048xbf16> loc(#loc1492)
      %32550 = stablehlo.convert %32549 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1486)
      %32551 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1487)
      %32552 = stablehlo.reduce(%32550 init: %32551) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1487)
        stablehlo.return %33956 : tensor<f32> loc(#loc1487)
      } loc(#loc1487)
      %32553 = stablehlo.broadcast_in_dim %32552, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1488)
      %32554 = stablehlo.convert %32553 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1489)
      %32555 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32556 = stablehlo.broadcast_in_dim %32555, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1490)
      %32557 = stablehlo.divide %32554, %32556 : tensor<1x1x2048x1xbf16> loc(#loc1490)
      %32558 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32559 = stablehlo.broadcast_in_dim %32558, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1493)
      %32560 = stablehlo.add %32557, %32559 : tensor<1x1x2048x1xbf16> loc(#loc1493)
      %32561 = stablehlo.rsqrt %32560 : tensor<1x1x2048x1xbf16> loc(#loc1494)
      %32562 = stablehlo.broadcast_in_dim %32561, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1495)
      %32563 = stablehlo.multiply %32546, %32562 : tensor<1x1x2048x2048xbf16> loc(#loc1495)
      %32564 = func.call @_where_378(%30042, %iterArg_380, %iterArg_381) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32565 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32566 = stablehlo.broadcast_in_dim %32565, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1496)
      %32567 = stablehlo.add %32566, %32564 : tensor<1x2048xbf16> loc(#loc1496)
      %32568 = stablehlo.broadcast_in_dim %32567, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1497)
      %32569 = stablehlo.broadcast_in_dim %32568, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1498)
      %32570 = stablehlo.multiply %32563, %32569 : tensor<1x1x2048x2048xbf16> loc(#loc1498)
      %32571 = func.call @_where_379(%30042, %iterArg_382, %iterArg_383) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32572 = stablehlo.broadcast_in_dim %32571, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1499)
      %32573 = stablehlo.broadcast_in_dim %32572, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1500)
      %32574 = stablehlo.add %32570, %32573 : tensor<1x1x2048x2048xbf16> loc(#loc1500)
      %32575 = func.call @_where_380(%30042, %iterArg_384, %iterArg_385) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %32576 = func.call @_einsum_381(%32574, %32575) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1501)
      %32577 = arith.select %ptrue, %32576, %32576 : tensor<1x1x2048x8192xbf16>
      %32578 = func.call @_where_382(%30042, %iterArg_386, %iterArg_387) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %32579 = stablehlo.broadcast_in_dim %32578, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1503)
      %32580 = stablehlo.broadcast_in_dim %32579, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1504)
      %32581 = stablehlo.add %32577, %32580 : tensor<1x1x2048x8192xbf16> loc(#loc1504)
      %32582 = stablehlo.multiply %32581, %32581 : tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32583 = stablehlo.multiply %32581, %32582 : tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32584 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %32585 = stablehlo.broadcast_in_dim %32584, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32586 = stablehlo.multiply %32585, %32583 : tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32587 = stablehlo.add %32581, %32586 : tensor<1x1x2048x8192xbf16> loc(#loc1506)
      %32588 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %32589 = stablehlo.broadcast_in_dim %32588, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32590 = stablehlo.multiply %32589, %32587 : tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32591 = stablehlo.tanh %32590 : tensor<1x1x2048x8192xbf16> loc(#loc1507)
      %32592 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32593 = stablehlo.broadcast_in_dim %32592, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1506)
      %32594 = stablehlo.add %32593, %32591 : tensor<1x1x2048x8192xbf16> loc(#loc1506)
      %32595 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %32596 = stablehlo.broadcast_in_dim %32595, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32597 = stablehlo.multiply %32596, %32594 : tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32598 = stablehlo.multiply %32581, %32597 : tensor<1x1x2048x8192xbf16> loc(#loc1505)
      %32599 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1508)
      %32600 = stablehlo.convert %32599 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1509)
      %32601 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32602 = stablehlo.broadcast_in_dim %32601, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1510)
      %32603 = stablehlo.subtract %32602, %32600 : tensor<1x1x2048x1xf32> loc(#loc1510)
      %32604 = stablehlo.convert %32603 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1511)
      %32605 = stablehlo.broadcast_in_dim %32604, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1512)
      %32606 = stablehlo.multiply %32598, %32605 : tensor<1x1x2048x8192xbf16> loc(#loc1512)
      %32607 = func.call @_where_383(%30042, %iterArg_388, %iterArg_389) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %32608 = func.call @_einsum_384(%32606, %32607) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1513)
      %32609 = arith.select %ptrue, %32608, %32608 : tensor<1x1x2048x2048xbf16>
      %32610 = func.call @_where_385(%30042, %iterArg_390, %iterArg_391) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32611 = stablehlo.broadcast_in_dim %32610, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1515)
      %32612 = stablehlo.broadcast_in_dim %32611, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1516)
      %32613 = stablehlo.add %32609, %32612 : tensor<1x1x2048x2048xbf16> loc(#loc1516)
      %32614 = stablehlo.convert %32599 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1509)
      %32615 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32616 = stablehlo.broadcast_in_dim %32615, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1510)
      %32617 = stablehlo.subtract %32616, %32614 : tensor<1x1x2048x1xf32> loc(#loc1510)
      %32618 = stablehlo.convert %32617 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1517)
      %32619 = stablehlo.broadcast_in_dim %32618, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1518)
      %32620 = stablehlo.multiply %32613, %32619 : tensor<1x1x2048x2048xbf16> loc(#loc1518)
      %32621 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32622 = stablehlo.broadcast_in_dim %32621, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1519)
      %32623 = stablehlo.multiply %32620, %32622 : tensor<1x1x2048x2048xbf16> loc(#loc1519)
      %32624 = stablehlo.add %32536, %32623 : tensor<1x1x2048x2048xbf16> loc(#loc1520)
      %32625 = stablehlo.convert %32624 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1521)
      %32626 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1522)
      %32627 = stablehlo.reduce(%32625 init: %32626) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1522)
        stablehlo.return %33956 : tensor<f32> loc(#loc1522)
      } loc(#loc1522)
      %32628 = stablehlo.broadcast_in_dim %32627, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1523)
      %32629 = stablehlo.convert %32628 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1524)
      %32630 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32631 = stablehlo.broadcast_in_dim %32630, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1525)
      %32632 = stablehlo.divide %32629, %32631 : tensor<1x1x2048x1xbf16> loc(#loc1525)
      %32633 = stablehlo.broadcast_in_dim %32632, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1526)
      %32634 = stablehlo.subtract %32624, %32633 : tensor<1x1x2048x2048xbf16> loc(#loc1526)
      %32635 = stablehlo.broadcast_in_dim %32632, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1526)
      %32636 = stablehlo.subtract %32624, %32635 : tensor<1x1x2048x2048xbf16> loc(#loc1526)
      %32637 = stablehlo.multiply %32636, %32636 : tensor<1x1x2048x2048xbf16> loc(#loc1527)
      %32638 = stablehlo.convert %32637 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1521)
      %32639 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1522)
      %32640 = stablehlo.reduce(%32638 init: %32639) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1522)
        stablehlo.return %33956 : tensor<f32> loc(#loc1522)
      } loc(#loc1522)
      %32641 = stablehlo.broadcast_in_dim %32640, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1523)
      %32642 = stablehlo.convert %32641 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1524)
      %32643 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32644 = stablehlo.broadcast_in_dim %32643, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1525)
      %32645 = stablehlo.divide %32642, %32644 : tensor<1x1x2048x1xbf16> loc(#loc1525)
      %32646 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32647 = stablehlo.broadcast_in_dim %32646, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1528)
      %32648 = stablehlo.add %32645, %32647 : tensor<1x1x2048x1xbf16> loc(#loc1528)
      %32649 = stablehlo.rsqrt %32648 : tensor<1x1x2048x1xbf16> loc(#loc1529)
      %32650 = stablehlo.broadcast_in_dim %32649, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1530)
      %32651 = stablehlo.multiply %32634, %32650 : tensor<1x1x2048x2048xbf16> loc(#loc1530)
      %32652 = func.call @_where_386(%30042, %iterArg_392, %iterArg_393) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32653 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32654 = stablehlo.broadcast_in_dim %32653, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1531)
      %32655 = stablehlo.add %32654, %32652 : tensor<1x2048xbf16> loc(#loc1531)
      %32656 = stablehlo.broadcast_in_dim %32655, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1532)
      %32657 = stablehlo.broadcast_in_dim %32656, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1533)
      %32658 = stablehlo.multiply %32651, %32657 : tensor<1x1x2048x2048xbf16> loc(#loc1533)
      %32659 = func.call @_where_387(%30042, %iterArg_394, %iterArg_395) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32660 = stablehlo.broadcast_in_dim %32659, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1534)
      %32661 = stablehlo.broadcast_in_dim %32660, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1535)
      %32662 = stablehlo.add %32658, %32661 : tensor<1x1x2048x2048xbf16> loc(#loc1535)
      %32663 = func.call @_where_388(%30042, %iterArg_396, %iterArg_397) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc372)
      %32664 = func.call @_einsum_389(%32662, %32663) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc1536)
      %32665 = "stablehlo.slice"(%32664) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1537)
      %32666 = stablehlo.reshape %32665 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1538)
      %32667 = arith.select %ptrue, %32666, %32666 : tensor<1x1x2048x32x64xbf16>
      %32668 = func.call @_where_390(%30042, %iterArg_398, %iterArg_399) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc372)
      %32669 = func.call @softplus_391(%32668, %iterArg_400) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1540)
      %32670 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc302)
      %32671 = stablehlo.broadcast_in_dim %32670, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1541)
      %32672 = stablehlo.multiply %32671, %32669 : tensor<1x64xbf16> loc(#loc1541)
      %32673 = stablehlo.broadcast_in_dim %32672, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc1542)
      %32674 = stablehlo.broadcast_in_dim %32673, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1543)
      %32675 = stablehlo.multiply %32667, %32674 : tensor<1x1x2048x32x64xbf16> loc(#loc1543)
      %32676 = "stablehlo.slice"(%32664) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1544)
      %32677 = stablehlo.reshape %32676 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1538)
      %32678 = arith.select %ptrue, %32677, %32677 : tensor<1x1x2048x32x64xbf16>
      %32679 = func.call @_einsum_393(%32675, %32678) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1545)
      %32680 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32681 = stablehlo.broadcast_in_dim %32680, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1546)
      %32682 = stablehlo.divide %32679, %32681 : tensor<1x1x32x2048x2048xbf16> loc(#loc1546)
      %32683 = stablehlo.tanh %32682 : tensor<1x1x32x2048x2048xbf16> loc(#loc1547)
      %32684 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc302)
      %32685 = stablehlo.broadcast_in_dim %32684, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1548)
      %32686 = stablehlo.multiply %32685, %32683 : tensor<1x1x32x2048x2048xbf16> loc(#loc1548)
      %32687 = stablehlo.convert %32686 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1549)
      %32688 = func.call @_where_394(%29952, %32687, %iterArg_401) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1550)
      %32689 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1551)
      %32690 = stablehlo.reduce(%32688 init: %32689) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1551)
        stablehlo.return %33956 : tensor<f32> loc(#loc1551)
      } loc(#loc1551)
      %32691 = stablehlo.broadcast_in_dim %32690, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1552)
      %32692 = stablehlo.broadcast_in_dim %32691, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1553)
      %32693 = stablehlo.subtract %32688, %32692 : tensor<1x1x32x2048x2048xf32> loc(#loc1553)
      %32694 = stablehlo.exponential %32693 : tensor<1x1x32x2048x2048xf32> loc(#loc1554)
      %32695 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1555)
      %32696 = stablehlo.reduce(%32694 init: %32695) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1555)
        stablehlo.return %33956 : tensor<f32> loc(#loc1555)
      } loc(#loc1555)
      %32697 = stablehlo.broadcast_in_dim %32696, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc1552)
      %32698 = stablehlo.broadcast_in_dim %32697, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc1556)
      %32699 = stablehlo.divide %32694, %32698 : tensor<1x1x32x2048x2048xf32> loc(#loc1556)
      %32700 = stablehlo.convert %32699 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc1557)
      %32701 = "stablehlo.slice"(%32664) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc1558)
      %32702 = stablehlo.reshape %32701 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1538)
      %32703 = arith.select %ptrue, %32702, %32702 : tensor<1x1x2048x32x64xbf16>
      %32704 = func.call @_einsum_395(%32700, %32703) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc1559)
      %32705 = arith.select %ptrue, %32704, %32704 : tensor<1x1x2048x32x64xbf16>
      %32706 = func.call @_where_396(%30042, %iterArg_402, %iterArg_403) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc372)
      %32707 = func.call @_einsum_397(%32705, %32706) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1560)
      %32708 = arith.select %ptrue, %32707, %32707 : tensor<1x1x2048x2048xbf16>
      %32709 = stablehlo.add %32708, %32624 : tensor<1x1x2048x2048xbf16> loc(#loc1562)
      %32710 = stablehlo.convert %32709 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1563)
      %32711 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1564)
      %32712 = stablehlo.reduce(%32710 init: %32711) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1564)
        stablehlo.return %33956 : tensor<f32> loc(#loc1564)
      } loc(#loc1564)
      %32713 = stablehlo.broadcast_in_dim %32712, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1565)
      %32714 = stablehlo.convert %32713 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1566)
      %32715 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32716 = stablehlo.broadcast_in_dim %32715, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1567)
      %32717 = stablehlo.divide %32714, %32716 : tensor<1x1x2048x1xbf16> loc(#loc1567)
      %32718 = stablehlo.broadcast_in_dim %32717, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1568)
      %32719 = stablehlo.subtract %32709, %32718 : tensor<1x1x2048x2048xbf16> loc(#loc1568)
      %32720 = stablehlo.broadcast_in_dim %32717, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1568)
      %32721 = stablehlo.subtract %32709, %32720 : tensor<1x1x2048x2048xbf16> loc(#loc1568)
      %32722 = stablehlo.multiply %32721, %32721 : tensor<1x1x2048x2048xbf16> loc(#loc1569)
      %32723 = stablehlo.convert %32722 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1563)
      %32724 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1564)
      %32725 = stablehlo.reduce(%32723 init: %32724) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1564)
        stablehlo.return %33956 : tensor<f32> loc(#loc1564)
      } loc(#loc1564)
      %32726 = stablehlo.broadcast_in_dim %32725, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc1565)
      %32727 = stablehlo.convert %32726 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1566)
      %32728 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc302)
      %32729 = stablehlo.broadcast_in_dim %32728, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1567)
      %32730 = stablehlo.divide %32727, %32729 : tensor<1x1x2048x1xbf16> loc(#loc1567)
      %32731 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc302)
      %32732 = stablehlo.broadcast_in_dim %32731, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc1570)
      %32733 = stablehlo.add %32730, %32732 : tensor<1x1x2048x1xbf16> loc(#loc1570)
      %32734 = stablehlo.rsqrt %32733 : tensor<1x1x2048x1xbf16> loc(#loc1571)
      %32735 = stablehlo.broadcast_in_dim %32734, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1572)
      %32736 = stablehlo.multiply %32719, %32735 : tensor<1x1x2048x2048xbf16> loc(#loc1572)
      %32737 = func.call @_where_398(%30042, %iterArg_404, %iterArg_405) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32738 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32739 = stablehlo.broadcast_in_dim %32738, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1573)
      %32740 = stablehlo.add %32739, %32737 : tensor<1x2048xbf16> loc(#loc1573)
      %32741 = stablehlo.broadcast_in_dim %32740, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1574)
      %32742 = stablehlo.broadcast_in_dim %32741, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1575)
      %32743 = stablehlo.multiply %32736, %32742 : tensor<1x1x2048x2048xbf16> loc(#loc1575)
      %32744 = func.call @_where_399(%30042, %iterArg_406, %iterArg_407) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32745 = stablehlo.broadcast_in_dim %32744, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1576)
      %32746 = stablehlo.broadcast_in_dim %32745, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1577)
      %32747 = stablehlo.add %32743, %32746 : tensor<1x1x2048x2048xbf16> loc(#loc1577)
      %32748 = func.call @_where_400(%30042, %iterArg_408, %iterArg_409) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc372)
      %32749 = func.call @_einsum_401(%32747, %32748) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1578)
      %32750 = arith.select %ptrue, %32749, %32749 : tensor<1x1x2048x8192xbf16>
      %32751 = func.call @_where_402(%30042, %iterArg_410, %iterArg_411) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc372)
      %32752 = stablehlo.broadcast_in_dim %32751, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc1580)
      %32753 = stablehlo.broadcast_in_dim %32752, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1581)
      %32754 = stablehlo.add %32750, %32753 : tensor<1x1x2048x8192xbf16> loc(#loc1581)
      %32755 = stablehlo.multiply %32754, %32754 : tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32756 = stablehlo.multiply %32754, %32755 : tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32757 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc302)
      %32758 = stablehlo.broadcast_in_dim %32757, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32759 = stablehlo.multiply %32758, %32756 : tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32760 = stablehlo.add %32754, %32759 : tensor<1x1x2048x8192xbf16> loc(#loc1583)
      %32761 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc302)
      %32762 = stablehlo.broadcast_in_dim %32761, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32763 = stablehlo.multiply %32762, %32760 : tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32764 = stablehlo.tanh %32763 : tensor<1x1x2048x8192xbf16> loc(#loc1584)
      %32765 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32766 = stablehlo.broadcast_in_dim %32765, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1583)
      %32767 = stablehlo.add %32766, %32764 : tensor<1x1x2048x8192xbf16> loc(#loc1583)
      %32768 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc302)
      %32769 = stablehlo.broadcast_in_dim %32768, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32770 = stablehlo.multiply %32769, %32767 : tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32771 = stablehlo.multiply %32754, %32770 : tensor<1x1x2048x8192xbf16> loc(#loc1582)
      %32772 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1585)
      %32773 = stablehlo.convert %32772 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1586)
      %32774 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32775 = stablehlo.broadcast_in_dim %32774, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1587)
      %32776 = stablehlo.subtract %32775, %32773 : tensor<1x1x2048x1xf32> loc(#loc1587)
      %32777 = stablehlo.convert %32776 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1588)
      %32778 = stablehlo.broadcast_in_dim %32777, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc1589)
      %32779 = stablehlo.multiply %32771, %32778 : tensor<1x1x2048x8192xbf16> loc(#loc1589)
      %32780 = func.call @_where_403(%30042, %iterArg_412, %iterArg_413) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc372)
      %32781 = func.call @_einsum_404(%32779, %32780) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1590)
      %32782 = arith.select %ptrue, %32781, %32781 : tensor<1x1x2048x2048xbf16>
      %32783 = func.call @_where_405(%30042, %iterArg_414, %iterArg_415) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc372)
      %32784 = stablehlo.broadcast_in_dim %32783, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc1592)
      %32785 = stablehlo.broadcast_in_dim %32784, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1593)
      %32786 = stablehlo.add %32782, %32785 : tensor<1x1x2048x2048xbf16> loc(#loc1593)
      %32787 = stablehlo.convert %32772 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1586)
      %32788 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32789 = stablehlo.broadcast_in_dim %32788, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1587)
      %32790 = stablehlo.subtract %32789, %32787 : tensor<1x1x2048x1xf32> loc(#loc1587)
      %32791 = stablehlo.convert %32790 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc1594)
      %32792 = stablehlo.broadcast_in_dim %32791, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1595)
      %32793 = stablehlo.multiply %32786, %32792 : tensor<1x1x2048x2048xbf16> loc(#loc1595)
      %32794 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc302)
      %32795 = stablehlo.broadcast_in_dim %32794, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1596)
      %32796 = stablehlo.multiply %32793, %32795 : tensor<1x1x2048x2048xbf16> loc(#loc1596)
      %32797 = stablehlo.add %32709, %32796 : tensor<1x1x2048x2048xbf16> loc(#loc1597)
      %32798 = func.call @_pad_406(%32797, %iterArg_416, %iterArg_417, %iterArg_418, %iterArg_419, %iterArg_420, %iterArg_421, %iterArg_422, %iterArg_423) : (tensor<1x1x2048x2048xbf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc1598)
      %32799 = "stablehlo.slice"(%32798) {limit_indices = dense<[1, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1599)
      %32800 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %32801 = stablehlo.add %iterArg_490, %32800 : tensor<i32> loc(#loc1600)
      %32802 = stablehlo.convert %30007 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1601)
      %32803 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1602)
      %32804 = stablehlo.convert %32803 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1603)
      %32805 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32806 = stablehlo.broadcast_in_dim %32805, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1604)
      %32807 = stablehlo.subtract %32806, %32804 : tensor<1x1x2048x1xf32> loc(#loc1604)
      %32808 = stablehlo.convert %32807 : tensor<1x1x2048x1xf32> loc(#loc1605)
      %32809 = stablehlo.broadcast_in_dim %32808, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32810 = stablehlo.multiply %32802, %32809 : tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32811 = stablehlo.abs %32810 : tensor<1x1x2048x2048xf32> loc(#loc1607)
      %32812 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1608)
      %32813 = stablehlo.reduce(%32811 init: %32812) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1608)
        stablehlo.return %33956 : tensor<f32> loc(#loc1608)
      } loc(#loc1608)
      %32814 = func.call @_where_407(%30042, %32813, %iterArg_424) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32815 = stablehlo.convert %32807 : tensor<1x1x2048x1xf32> loc(#loc1605)
      %32816 = stablehlo.broadcast_in_dim %32815, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32817 = stablehlo.multiply %32802, %32816 : tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32818 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1610)
      %32819 = stablehlo.reduce(%32817 init: %32818) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1610)
        stablehlo.return %33956 : tensor<f32> loc(#loc1610)
      } loc(#loc1610)
      %32820 = stablehlo.convert %32807 : tensor<1x1x2048x1xf32> loc(#loc1605)
      %32821 = stablehlo.broadcast_in_dim %32820, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32822 = stablehlo.multiply %iterArg_425, %32821 : tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32823 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1610)
      %32824 = stablehlo.reduce(%32822 init: %32823) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1610)
        stablehlo.return %33956 : tensor<f32> loc(#loc1610)
      } loc(#loc1610)
      %32825 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32826 = stablehlo.broadcast_in_dim %32825, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1611)
      %32827 = stablehlo.maximum %32826, %32824 : tensor<1xf32> loc(#loc1611)
      %32828 = stablehlo.divide %32819, %32827 : tensor<1xf32> loc(#loc1612)
      %32829 = func.call @_where_408(%30042, %32828, %iterArg_426) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32830 = stablehlo.broadcast_in_dim %32828, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1613)
      %32831 = stablehlo.broadcast_in_dim %32830, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1614)
      %32832 = stablehlo.subtract %32802, %32831 : tensor<1x1x2048x2048xf32> loc(#loc1614)
      %32833 = stablehlo.convert %32807 : tensor<1x1x2048x1xf32> loc(#loc1605)
      %32834 = stablehlo.broadcast_in_dim %32833, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32835 = stablehlo.multiply %32832, %32834 : tensor<1x1x2048x2048xf32> loc(#loc1606)
      %32836 = stablehlo.multiply %32835, %32835 : tensor<1x1x2048x2048xf32> loc(#loc1615)
      %32837 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1610)
      %32838 = stablehlo.reduce(%32836 init: %32837) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1610)
        stablehlo.return %33956 : tensor<f32> loc(#loc1610)
      } loc(#loc1610)
      %32839 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32840 = stablehlo.broadcast_in_dim %32839, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1611)
      %32841 = stablehlo.maximum %32840, %32824 : tensor<1xf32> loc(#loc1611)
      %32842 = stablehlo.divide %32838, %32841 : tensor<1xf32> loc(#loc1612)
      %32843 = stablehlo.sqrt %32842 : tensor<1xf32> loc(#loc1616)
      %32844 = func.call @_where_409(%30042, %32843, %iterArg_427) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32845 = stablehlo.convert %30202 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1617)
      %32846 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1618)
      %32847 = stablehlo.convert %32846 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1619)
      %32848 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32849 = stablehlo.broadcast_in_dim %32848, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1620)
      %32850 = stablehlo.subtract %32849, %32847 : tensor<1x1x2048x1xf32> loc(#loc1620)
      %32851 = stablehlo.convert %32850 : tensor<1x1x2048x1xf32> loc(#loc1621)
      %32852 = stablehlo.broadcast_in_dim %32851, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32853 = stablehlo.multiply %32845, %32852 : tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32854 = stablehlo.abs %32853 : tensor<1x1x2048x2048xf32> loc(#loc1623)
      %32855 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1624)
      %32856 = stablehlo.reduce(%32854 init: %32855) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1624)
        stablehlo.return %33956 : tensor<f32> loc(#loc1624)
      } loc(#loc1624)
      %32857 = func.call @_where_410(%30042, %32856, %iterArg_428) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32858 = stablehlo.convert %32850 : tensor<1x1x2048x1xf32> loc(#loc1621)
      %32859 = stablehlo.broadcast_in_dim %32858, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32860 = stablehlo.multiply %32845, %32859 : tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32861 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1625)
      %32862 = stablehlo.reduce(%32860 init: %32861) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1625)
        stablehlo.return %33956 : tensor<f32> loc(#loc1625)
      } loc(#loc1625)
      %32863 = stablehlo.convert %32850 : tensor<1x1x2048x1xf32> loc(#loc1621)
      %32864 = stablehlo.broadcast_in_dim %32863, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32865 = stablehlo.multiply %iterArg_429, %32864 : tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32866 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1625)
      %32867 = stablehlo.reduce(%32865 init: %32866) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1625)
        stablehlo.return %33956 : tensor<f32> loc(#loc1625)
      } loc(#loc1625)
      %32868 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32869 = stablehlo.broadcast_in_dim %32868, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1626)
      %32870 = stablehlo.maximum %32869, %32867 : tensor<1xf32> loc(#loc1626)
      %32871 = stablehlo.divide %32862, %32870 : tensor<1xf32> loc(#loc1627)
      %32872 = func.call @_where_411(%30042, %32871, %iterArg_430) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32873 = stablehlo.broadcast_in_dim %32871, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1628)
      %32874 = stablehlo.broadcast_in_dim %32873, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1629)
      %32875 = stablehlo.subtract %32845, %32874 : tensor<1x1x2048x2048xf32> loc(#loc1629)
      %32876 = stablehlo.convert %32850 : tensor<1x1x2048x1xf32> loc(#loc1621)
      %32877 = stablehlo.broadcast_in_dim %32876, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32878 = stablehlo.multiply %32875, %32877 : tensor<1x1x2048x2048xf32> loc(#loc1622)
      %32879 = stablehlo.multiply %32878, %32878 : tensor<1x1x2048x2048xf32> loc(#loc1630)
      %32880 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1625)
      %32881 = stablehlo.reduce(%32879 init: %32880) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1625)
        stablehlo.return %33956 : tensor<f32> loc(#loc1625)
      } loc(#loc1625)
      %32882 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32883 = stablehlo.broadcast_in_dim %32882, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1626)
      %32884 = stablehlo.maximum %32883, %32867 : tensor<1xf32> loc(#loc1626)
      %32885 = stablehlo.divide %32881, %32884 : tensor<1xf32> loc(#loc1627)
      %32886 = stablehlo.sqrt %32885 : tensor<1xf32> loc(#loc1631)
      %32887 = func.call @_where_412(%30042, %32886, %iterArg_431) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32888 = stablehlo.convert %31759 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1632)
      %32889 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1633)
      %32890 = stablehlo.convert %32889 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1634)
      %32891 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32892 = stablehlo.broadcast_in_dim %32891, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1635)
      %32893 = stablehlo.subtract %32892, %32890 : tensor<1x1x2048x1xf32> loc(#loc1635)
      %32894 = stablehlo.convert %32893 : tensor<1x1x2048x1xf32> loc(#loc1636)
      %32895 = stablehlo.broadcast_in_dim %32894, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32896 = stablehlo.multiply %32888, %32895 : tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32897 = stablehlo.abs %32896 : tensor<1x1x2048x2048xf32> loc(#loc1638)
      %32898 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1639)
      %32899 = stablehlo.reduce(%32897 init: %32898) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1639)
        stablehlo.return %33956 : tensor<f32> loc(#loc1639)
      } loc(#loc1639)
      %32900 = func.call @_where_413(%30042, %32899, %iterArg_432) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32901 = stablehlo.convert %32893 : tensor<1x1x2048x1xf32> loc(#loc1636)
      %32902 = stablehlo.broadcast_in_dim %32901, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32903 = stablehlo.multiply %32888, %32902 : tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32904 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1640)
      %32905 = stablehlo.reduce(%32903 init: %32904) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1640)
        stablehlo.return %33956 : tensor<f32> loc(#loc1640)
      } loc(#loc1640)
      %32906 = stablehlo.convert %32893 : tensor<1x1x2048x1xf32> loc(#loc1636)
      %32907 = stablehlo.broadcast_in_dim %32906, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32908 = stablehlo.multiply %iterArg_433, %32907 : tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32909 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1640)
      %32910 = stablehlo.reduce(%32908 init: %32909) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1640)
        stablehlo.return %33956 : tensor<f32> loc(#loc1640)
      } loc(#loc1640)
      %32911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32912 = stablehlo.broadcast_in_dim %32911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1641)
      %32913 = stablehlo.maximum %32912, %32910 : tensor<1xf32> loc(#loc1641)
      %32914 = stablehlo.divide %32905, %32913 : tensor<1xf32> loc(#loc1642)
      %32915 = func.call @_where_414(%30042, %32914, %iterArg_434) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32916 = stablehlo.broadcast_in_dim %32914, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1643)
      %32917 = stablehlo.broadcast_in_dim %32916, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1644)
      %32918 = stablehlo.subtract %32888, %32917 : tensor<1x1x2048x2048xf32> loc(#loc1644)
      %32919 = stablehlo.convert %32893 : tensor<1x1x2048x1xf32> loc(#loc1636)
      %32920 = stablehlo.broadcast_in_dim %32919, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32921 = stablehlo.multiply %32918, %32920 : tensor<1x1x2048x2048xf32> loc(#loc1637)
      %32922 = stablehlo.multiply %32921, %32921 : tensor<1x1x2048x2048xf32> loc(#loc1645)
      %32923 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1640)
      %32924 = stablehlo.reduce(%32922 init: %32923) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1640)
        stablehlo.return %33956 : tensor<f32> loc(#loc1640)
      } loc(#loc1640)
      %32925 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32926 = stablehlo.broadcast_in_dim %32925, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1641)
      %32927 = stablehlo.maximum %32926, %32910 : tensor<1xf32> loc(#loc1641)
      %32928 = stablehlo.divide %32924, %32927 : tensor<1xf32> loc(#loc1642)
      %32929 = stablehlo.sqrt %32928 : tensor<1xf32> loc(#loc1646)
      %32930 = func.call @_where_415(%30042, %32929, %iterArg_435) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32931 = stablehlo.convert %31932 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1647)
      %32932 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1648)
      %32933 = stablehlo.convert %32932 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1649)
      %32934 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32935 = stablehlo.broadcast_in_dim %32934, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1650)
      %32936 = stablehlo.subtract %32935, %32933 : tensor<1x1x2048x1xf32> loc(#loc1650)
      %32937 = stablehlo.convert %32936 : tensor<1x1x2048x1xf32> loc(#loc1651)
      %32938 = stablehlo.broadcast_in_dim %32937, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32939 = stablehlo.multiply %32931, %32938 : tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32940 = stablehlo.abs %32939 : tensor<1x1x2048x2048xf32> loc(#loc1653)
      %32941 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1654)
      %32942 = stablehlo.reduce(%32940 init: %32941) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1654)
        stablehlo.return %33956 : tensor<f32> loc(#loc1654)
      } loc(#loc1654)
      %32943 = func.call @_where_416(%30042, %32942, %iterArg_436) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32944 = stablehlo.convert %32936 : tensor<1x1x2048x1xf32> loc(#loc1651)
      %32945 = stablehlo.broadcast_in_dim %32944, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32946 = stablehlo.multiply %32931, %32945 : tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32947 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1655)
      %32948 = stablehlo.reduce(%32946 init: %32947) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1655)
        stablehlo.return %33956 : tensor<f32> loc(#loc1655)
      } loc(#loc1655)
      %32949 = stablehlo.convert %32936 : tensor<1x1x2048x1xf32> loc(#loc1651)
      %32950 = stablehlo.broadcast_in_dim %32949, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32951 = stablehlo.multiply %iterArg_437, %32950 : tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32952 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1655)
      %32953 = stablehlo.reduce(%32951 init: %32952) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1655)
        stablehlo.return %33956 : tensor<f32> loc(#loc1655)
      } loc(#loc1655)
      %32954 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32955 = stablehlo.broadcast_in_dim %32954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1656)
      %32956 = stablehlo.maximum %32955, %32953 : tensor<1xf32> loc(#loc1656)
      %32957 = stablehlo.divide %32948, %32956 : tensor<1xf32> loc(#loc1657)
      %32958 = func.call @_where_417(%30042, %32957, %iterArg_438) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32959 = stablehlo.broadcast_in_dim %32957, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1658)
      %32960 = stablehlo.broadcast_in_dim %32959, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1659)
      %32961 = stablehlo.subtract %32931, %32960 : tensor<1x1x2048x2048xf32> loc(#loc1659)
      %32962 = stablehlo.convert %32936 : tensor<1x1x2048x1xf32> loc(#loc1651)
      %32963 = stablehlo.broadcast_in_dim %32962, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32964 = stablehlo.multiply %32961, %32963 : tensor<1x1x2048x2048xf32> loc(#loc1652)
      %32965 = stablehlo.multiply %32964, %32964 : tensor<1x1x2048x2048xf32> loc(#loc1660)
      %32966 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1655)
      %32967 = stablehlo.reduce(%32965 init: %32966) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1655)
        stablehlo.return %33956 : tensor<f32> loc(#loc1655)
      } loc(#loc1655)
      %32968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32969 = stablehlo.broadcast_in_dim %32968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1656)
      %32970 = stablehlo.maximum %32969, %32953 : tensor<1xf32> loc(#loc1656)
      %32971 = stablehlo.divide %32967, %32970 : tensor<1xf32> loc(#loc1657)
      %32972 = stablehlo.sqrt %32971 : tensor<1xf32> loc(#loc1661)
      %32973 = func.call @_where_418(%30042, %32972, %iterArg_439) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32974 = stablehlo.convert %32105 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1662)
      %32975 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1663)
      %32976 = stablehlo.convert %32975 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1664)
      %32977 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32978 = stablehlo.broadcast_in_dim %32977, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1665)
      %32979 = stablehlo.subtract %32978, %32976 : tensor<1x1x2048x1xf32> loc(#loc1665)
      %32980 = stablehlo.convert %32979 : tensor<1x1x2048x1xf32> loc(#loc1666)
      %32981 = stablehlo.broadcast_in_dim %32980, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1667)
      %32982 = stablehlo.multiply %32974, %32981 : tensor<1x1x2048x2048xf32> loc(#loc1667)
      %32983 = stablehlo.abs %32982 : tensor<1x1x2048x2048xf32> loc(#loc1668)
      %32984 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1669)
      %32985 = stablehlo.reduce(%32983 init: %32984) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1669)
        stablehlo.return %33956 : tensor<f32> loc(#loc1669)
      } loc(#loc1669)
      %32986 = func.call @_where_419(%30042, %32985, %iterArg_440) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %32987 = stablehlo.convert %32979 : tensor<1x1x2048x1xf32> loc(#loc1666)
      %32988 = stablehlo.broadcast_in_dim %32987, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1667)
      %32989 = stablehlo.multiply %32974, %32988 : tensor<1x1x2048x2048xf32> loc(#loc1667)
      %32990 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1670)
      %32991 = stablehlo.reduce(%32989 init: %32990) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1670)
        stablehlo.return %33956 : tensor<f32> loc(#loc1670)
      } loc(#loc1670)
      %32992 = stablehlo.convert %32979 : tensor<1x1x2048x1xf32> loc(#loc1666)
      %32993 = stablehlo.broadcast_in_dim %32992, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1667)
      %32994 = stablehlo.multiply %iterArg_441, %32993 : tensor<1x1x2048x2048xf32> loc(#loc1667)
      %32995 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1670)
      %32996 = stablehlo.reduce(%32994 init: %32995) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1670)
        stablehlo.return %33956 : tensor<f32> loc(#loc1670)
      } loc(#loc1670)
      %32997 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %32998 = stablehlo.broadcast_in_dim %32997, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1671)
      %32999 = stablehlo.maximum %32998, %32996 : tensor<1xf32> loc(#loc1671)
      %33000 = stablehlo.divide %32991, %32999 : tensor<1xf32> loc(#loc1672)
      %33001 = func.call @_where_420(%30042, %33000, %iterArg_442) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33002 = stablehlo.broadcast_in_dim %33000, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1673)
      %33003 = stablehlo.broadcast_in_dim %33002, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1674)
      %33004 = stablehlo.subtract %32974, %33003 : tensor<1x1x2048x2048xf32> loc(#loc1674)
      %33005 = stablehlo.convert %32979 : tensor<1x1x2048x1xf32> loc(#loc1666)
      %33006 = stablehlo.broadcast_in_dim %33005, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1667)
      %33007 = stablehlo.multiply %33004, %33006 : tensor<1x1x2048x2048xf32> loc(#loc1667)
      %33008 = stablehlo.multiply %33007, %33007 : tensor<1x1x2048x2048xf32> loc(#loc1675)
      %33009 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1670)
      %33010 = stablehlo.reduce(%33008 init: %33009) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1670)
        stablehlo.return %33956 : tensor<f32> loc(#loc1670)
      } loc(#loc1670)
      %33011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33012 = stablehlo.broadcast_in_dim %33011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1671)
      %33013 = stablehlo.maximum %33012, %32996 : tensor<1xf32> loc(#loc1671)
      %33014 = stablehlo.divide %33010, %33013 : tensor<1xf32> loc(#loc1672)
      %33015 = stablehlo.sqrt %33014 : tensor<1xf32> loc(#loc1676)
      %33016 = func.call @_where_421(%30042, %33015, %iterArg_443) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33017 = stablehlo.convert %32278 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1677)
      %33018 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1678)
      %33019 = stablehlo.convert %33018 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1679)
      %33020 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33021 = stablehlo.broadcast_in_dim %33020, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1680)
      %33022 = stablehlo.subtract %33021, %33019 : tensor<1x1x2048x1xf32> loc(#loc1680)
      %33023 = stablehlo.convert %33022 : tensor<1x1x2048x1xf32> loc(#loc1681)
      %33024 = stablehlo.broadcast_in_dim %33023, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33025 = stablehlo.multiply %33017, %33024 : tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33026 = stablehlo.abs %33025 : tensor<1x1x2048x2048xf32> loc(#loc1683)
      %33027 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1684)
      %33028 = stablehlo.reduce(%33026 init: %33027) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1684)
        stablehlo.return %33956 : tensor<f32> loc(#loc1684)
      } loc(#loc1684)
      %33029 = func.call @_where_422(%30042, %33028, %iterArg_444) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33030 = stablehlo.convert %33022 : tensor<1x1x2048x1xf32> loc(#loc1681)
      %33031 = stablehlo.broadcast_in_dim %33030, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33032 = stablehlo.multiply %33017, %33031 : tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33033 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1685)
      %33034 = stablehlo.reduce(%33032 init: %33033) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1685)
        stablehlo.return %33956 : tensor<f32> loc(#loc1685)
      } loc(#loc1685)
      %33035 = stablehlo.convert %33022 : tensor<1x1x2048x1xf32> loc(#loc1681)
      %33036 = stablehlo.broadcast_in_dim %33035, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33037 = stablehlo.multiply %iterArg_445, %33036 : tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33038 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1685)
      %33039 = stablehlo.reduce(%33037 init: %33038) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1685)
        stablehlo.return %33956 : tensor<f32> loc(#loc1685)
      } loc(#loc1685)
      %33040 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33041 = stablehlo.broadcast_in_dim %33040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1686)
      %33042 = stablehlo.maximum %33041, %33039 : tensor<1xf32> loc(#loc1686)
      %33043 = stablehlo.divide %33034, %33042 : tensor<1xf32> loc(#loc1687)
      %33044 = func.call @_where_423(%30042, %33043, %iterArg_446) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33045 = stablehlo.broadcast_in_dim %33043, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1688)
      %33046 = stablehlo.broadcast_in_dim %33045, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1689)
      %33047 = stablehlo.subtract %33017, %33046 : tensor<1x1x2048x2048xf32> loc(#loc1689)
      %33048 = stablehlo.convert %33022 : tensor<1x1x2048x1xf32> loc(#loc1681)
      %33049 = stablehlo.broadcast_in_dim %33048, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33050 = stablehlo.multiply %33047, %33049 : tensor<1x1x2048x2048xf32> loc(#loc1682)
      %33051 = stablehlo.multiply %33050, %33050 : tensor<1x1x2048x2048xf32> loc(#loc1690)
      %33052 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1685)
      %33053 = stablehlo.reduce(%33051 init: %33052) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1685)
        stablehlo.return %33956 : tensor<f32> loc(#loc1685)
      } loc(#loc1685)
      %33054 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33055 = stablehlo.broadcast_in_dim %33054, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1686)
      %33056 = stablehlo.maximum %33055, %33039 : tensor<1xf32> loc(#loc1686)
      %33057 = stablehlo.divide %33053, %33056 : tensor<1xf32> loc(#loc1687)
      %33058 = stablehlo.sqrt %33057 : tensor<1xf32> loc(#loc1691)
      %33059 = func.call @_where_424(%30042, %33058, %iterArg_447) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33060 = stablehlo.convert %32451 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1692)
      %33061 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1693)
      %33062 = stablehlo.convert %33061 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1694)
      %33063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33064 = stablehlo.broadcast_in_dim %33063, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1695)
      %33065 = stablehlo.subtract %33064, %33062 : tensor<1x1x2048x1xf32> loc(#loc1695)
      %33066 = stablehlo.convert %33065 : tensor<1x1x2048x1xf32> loc(#loc1696)
      %33067 = stablehlo.broadcast_in_dim %33066, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33068 = stablehlo.multiply %33060, %33067 : tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33069 = stablehlo.abs %33068 : tensor<1x1x2048x2048xf32> loc(#loc1698)
      %33070 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1699)
      %33071 = stablehlo.reduce(%33069 init: %33070) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1699)
        stablehlo.return %33956 : tensor<f32> loc(#loc1699)
      } loc(#loc1699)
      %33072 = func.call @_where_425(%30042, %33071, %iterArg_448) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33073 = stablehlo.convert %33065 : tensor<1x1x2048x1xf32> loc(#loc1696)
      %33074 = stablehlo.broadcast_in_dim %33073, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33075 = stablehlo.multiply %33060, %33074 : tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33076 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1700)
      %33077 = stablehlo.reduce(%33075 init: %33076) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1700)
        stablehlo.return %33956 : tensor<f32> loc(#loc1700)
      } loc(#loc1700)
      %33078 = stablehlo.convert %33065 : tensor<1x1x2048x1xf32> loc(#loc1696)
      %33079 = stablehlo.broadcast_in_dim %33078, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33080 = stablehlo.multiply %iterArg_449, %33079 : tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33081 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1700)
      %33082 = stablehlo.reduce(%33080 init: %33081) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1700)
        stablehlo.return %33956 : tensor<f32> loc(#loc1700)
      } loc(#loc1700)
      %33083 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33084 = stablehlo.broadcast_in_dim %33083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1701)
      %33085 = stablehlo.maximum %33084, %33082 : tensor<1xf32> loc(#loc1701)
      %33086 = stablehlo.divide %33077, %33085 : tensor<1xf32> loc(#loc1702)
      %33087 = func.call @_where_426(%30042, %33086, %iterArg_450) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33088 = stablehlo.broadcast_in_dim %33086, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1703)
      %33089 = stablehlo.broadcast_in_dim %33088, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1704)
      %33090 = stablehlo.subtract %33060, %33089 : tensor<1x1x2048x2048xf32> loc(#loc1704)
      %33091 = stablehlo.convert %33065 : tensor<1x1x2048x1xf32> loc(#loc1696)
      %33092 = stablehlo.broadcast_in_dim %33091, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33093 = stablehlo.multiply %33090, %33092 : tensor<1x1x2048x2048xf32> loc(#loc1697)
      %33094 = stablehlo.multiply %33093, %33093 : tensor<1x1x2048x2048xf32> loc(#loc1705)
      %33095 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1700)
      %33096 = stablehlo.reduce(%33094 init: %33095) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1700)
        stablehlo.return %33956 : tensor<f32> loc(#loc1700)
      } loc(#loc1700)
      %33097 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33098 = stablehlo.broadcast_in_dim %33097, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1701)
      %33099 = stablehlo.maximum %33098, %33082 : tensor<1xf32> loc(#loc1701)
      %33100 = stablehlo.divide %33096, %33099 : tensor<1xf32> loc(#loc1702)
      %33101 = stablehlo.sqrt %33100 : tensor<1xf32> loc(#loc1706)
      %33102 = func.call @_where_427(%30042, %33101, %iterArg_451) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33103 = stablehlo.convert %32624 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1707)
      %33104 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1708)
      %33105 = stablehlo.convert %33104 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1709)
      %33106 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33107 = stablehlo.broadcast_in_dim %33106, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1710)
      %33108 = stablehlo.subtract %33107, %33105 : tensor<1x1x2048x1xf32> loc(#loc1710)
      %33109 = stablehlo.convert %33108 : tensor<1x1x2048x1xf32> loc(#loc1711)
      %33110 = stablehlo.broadcast_in_dim %33109, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33111 = stablehlo.multiply %33103, %33110 : tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33112 = stablehlo.abs %33111 : tensor<1x1x2048x2048xf32> loc(#loc1713)
      %33113 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1714)
      %33114 = stablehlo.reduce(%33112 init: %33113) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1714)
        stablehlo.return %33956 : tensor<f32> loc(#loc1714)
      } loc(#loc1714)
      %33115 = func.call @_where_428(%30042, %33114, %iterArg_452) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33116 = stablehlo.convert %33108 : tensor<1x1x2048x1xf32> loc(#loc1711)
      %33117 = stablehlo.broadcast_in_dim %33116, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33118 = stablehlo.multiply %33103, %33117 : tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33119 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1715)
      %33120 = stablehlo.reduce(%33118 init: %33119) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1715)
        stablehlo.return %33956 : tensor<f32> loc(#loc1715)
      } loc(#loc1715)
      %33121 = stablehlo.convert %33108 : tensor<1x1x2048x1xf32> loc(#loc1711)
      %33122 = stablehlo.broadcast_in_dim %33121, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33123 = stablehlo.multiply %iterArg_453, %33122 : tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33124 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1715)
      %33125 = stablehlo.reduce(%33123 init: %33124) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1715)
        stablehlo.return %33956 : tensor<f32> loc(#loc1715)
      } loc(#loc1715)
      %33126 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33127 = stablehlo.broadcast_in_dim %33126, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1716)
      %33128 = stablehlo.maximum %33127, %33125 : tensor<1xf32> loc(#loc1716)
      %33129 = stablehlo.divide %33120, %33128 : tensor<1xf32> loc(#loc1717)
      %33130 = func.call @_where_429(%30042, %33129, %iterArg_454) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33131 = stablehlo.broadcast_in_dim %33129, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1718)
      %33132 = stablehlo.broadcast_in_dim %33131, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1719)
      %33133 = stablehlo.subtract %33103, %33132 : tensor<1x1x2048x2048xf32> loc(#loc1719)
      %33134 = stablehlo.convert %33108 : tensor<1x1x2048x1xf32> loc(#loc1711)
      %33135 = stablehlo.broadcast_in_dim %33134, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33136 = stablehlo.multiply %33133, %33135 : tensor<1x1x2048x2048xf32> loc(#loc1712)
      %33137 = stablehlo.multiply %33136, %33136 : tensor<1x1x2048x2048xf32> loc(#loc1720)
      %33138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1715)
      %33139 = stablehlo.reduce(%33137 init: %33138) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1715)
        stablehlo.return %33956 : tensor<f32> loc(#loc1715)
      } loc(#loc1715)
      %33140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33141 = stablehlo.broadcast_in_dim %33140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1716)
      %33142 = stablehlo.maximum %33141, %33125 : tensor<1xf32> loc(#loc1716)
      %33143 = stablehlo.divide %33139, %33142 : tensor<1xf32> loc(#loc1717)
      %33144 = stablehlo.sqrt %33143 : tensor<1xf32> loc(#loc1721)
      %33145 = func.call @_where_430(%30042, %33144, %iterArg_455) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33146 = stablehlo.convert %30375 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1722)
      %33147 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1723)
      %33148 = stablehlo.convert %33147 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1724)
      %33149 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33150 = stablehlo.broadcast_in_dim %33149, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1725)
      %33151 = stablehlo.subtract %33150, %33148 : tensor<1x1x2048x1xf32> loc(#loc1725)
      %33152 = stablehlo.convert %33151 : tensor<1x1x2048x1xf32> loc(#loc1726)
      %33153 = stablehlo.broadcast_in_dim %33152, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33154 = stablehlo.multiply %33146, %33153 : tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33155 = stablehlo.abs %33154 : tensor<1x1x2048x2048xf32> loc(#loc1728)
      %33156 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1729)
      %33157 = stablehlo.reduce(%33155 init: %33156) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1729)
        stablehlo.return %33956 : tensor<f32> loc(#loc1729)
      } loc(#loc1729)
      %33158 = func.call @_where_431(%30042, %33157, %iterArg_456) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33159 = stablehlo.convert %33151 : tensor<1x1x2048x1xf32> loc(#loc1726)
      %33160 = stablehlo.broadcast_in_dim %33159, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33161 = stablehlo.multiply %33146, %33160 : tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33162 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1730)
      %33163 = stablehlo.reduce(%33161 init: %33162) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1730)
        stablehlo.return %33956 : tensor<f32> loc(#loc1730)
      } loc(#loc1730)
      %33164 = stablehlo.convert %33151 : tensor<1x1x2048x1xf32> loc(#loc1726)
      %33165 = stablehlo.broadcast_in_dim %33164, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33166 = stablehlo.multiply %iterArg_457, %33165 : tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33167 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1730)
      %33168 = stablehlo.reduce(%33166 init: %33167) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1730)
        stablehlo.return %33956 : tensor<f32> loc(#loc1730)
      } loc(#loc1730)
      %33169 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33170 = stablehlo.broadcast_in_dim %33169, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1731)
      %33171 = stablehlo.maximum %33170, %33168 : tensor<1xf32> loc(#loc1731)
      %33172 = stablehlo.divide %33163, %33171 : tensor<1xf32> loc(#loc1732)
      %33173 = func.call @_where_432(%30042, %33172, %iterArg_458) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33174 = stablehlo.broadcast_in_dim %33172, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1733)
      %33175 = stablehlo.broadcast_in_dim %33174, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1734)
      %33176 = stablehlo.subtract %33146, %33175 : tensor<1x1x2048x2048xf32> loc(#loc1734)
      %33177 = stablehlo.convert %33151 : tensor<1x1x2048x1xf32> loc(#loc1726)
      %33178 = stablehlo.broadcast_in_dim %33177, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33179 = stablehlo.multiply %33176, %33178 : tensor<1x1x2048x2048xf32> loc(#loc1727)
      %33180 = stablehlo.multiply %33179, %33179 : tensor<1x1x2048x2048xf32> loc(#loc1735)
      %33181 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1730)
      %33182 = stablehlo.reduce(%33180 init: %33181) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1730)
        stablehlo.return %33956 : tensor<f32> loc(#loc1730)
      } loc(#loc1730)
      %33183 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33184 = stablehlo.broadcast_in_dim %33183, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1731)
      %33185 = stablehlo.maximum %33184, %33168 : tensor<1xf32> loc(#loc1731)
      %33186 = stablehlo.divide %33182, %33185 : tensor<1xf32> loc(#loc1732)
      %33187 = stablehlo.sqrt %33186 : tensor<1xf32> loc(#loc1736)
      %33188 = func.call @_where_433(%30042, %33187, %iterArg_459) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33189 = stablehlo.convert %30548 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1737)
      %33190 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1738)
      %33191 = stablehlo.convert %33190 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1739)
      %33192 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33193 = stablehlo.broadcast_in_dim %33192, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1740)
      %33194 = stablehlo.subtract %33193, %33191 : tensor<1x1x2048x1xf32> loc(#loc1740)
      %33195 = stablehlo.convert %33194 : tensor<1x1x2048x1xf32> loc(#loc1741)
      %33196 = stablehlo.broadcast_in_dim %33195, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33197 = stablehlo.multiply %33189, %33196 : tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33198 = stablehlo.abs %33197 : tensor<1x1x2048x2048xf32> loc(#loc1743)
      %33199 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1744)
      %33200 = stablehlo.reduce(%33198 init: %33199) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1744)
        stablehlo.return %33956 : tensor<f32> loc(#loc1744)
      } loc(#loc1744)
      %33201 = func.call @_where_434(%30042, %33200, %iterArg_460) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33202 = stablehlo.convert %33194 : tensor<1x1x2048x1xf32> loc(#loc1741)
      %33203 = stablehlo.broadcast_in_dim %33202, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33204 = stablehlo.multiply %33189, %33203 : tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33205 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1745)
      %33206 = stablehlo.reduce(%33204 init: %33205) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1745)
        stablehlo.return %33956 : tensor<f32> loc(#loc1745)
      } loc(#loc1745)
      %33207 = stablehlo.convert %33194 : tensor<1x1x2048x1xf32> loc(#loc1741)
      %33208 = stablehlo.broadcast_in_dim %33207, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33209 = stablehlo.multiply %iterArg_461, %33208 : tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33210 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1745)
      %33211 = stablehlo.reduce(%33209 init: %33210) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1745)
        stablehlo.return %33956 : tensor<f32> loc(#loc1745)
      } loc(#loc1745)
      %33212 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33213 = stablehlo.broadcast_in_dim %33212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1746)
      %33214 = stablehlo.maximum %33213, %33211 : tensor<1xf32> loc(#loc1746)
      %33215 = stablehlo.divide %33206, %33214 : tensor<1xf32> loc(#loc1747)
      %33216 = func.call @_where_435(%30042, %33215, %iterArg_462) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33217 = stablehlo.broadcast_in_dim %33215, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1748)
      %33218 = stablehlo.broadcast_in_dim %33217, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1749)
      %33219 = stablehlo.subtract %33189, %33218 : tensor<1x1x2048x2048xf32> loc(#loc1749)
      %33220 = stablehlo.convert %33194 : tensor<1x1x2048x1xf32> loc(#loc1741)
      %33221 = stablehlo.broadcast_in_dim %33220, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33222 = stablehlo.multiply %33219, %33221 : tensor<1x1x2048x2048xf32> loc(#loc1742)
      %33223 = stablehlo.multiply %33222, %33222 : tensor<1x1x2048x2048xf32> loc(#loc1750)
      %33224 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1745)
      %33225 = stablehlo.reduce(%33223 init: %33224) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1745)
        stablehlo.return %33956 : tensor<f32> loc(#loc1745)
      } loc(#loc1745)
      %33226 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33227 = stablehlo.broadcast_in_dim %33226, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1746)
      %33228 = stablehlo.maximum %33227, %33211 : tensor<1xf32> loc(#loc1746)
      %33229 = stablehlo.divide %33225, %33228 : tensor<1xf32> loc(#loc1747)
      %33230 = stablehlo.sqrt %33229 : tensor<1xf32> loc(#loc1751)
      %33231 = func.call @_where_436(%30042, %33230, %iterArg_463) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33232 = stablehlo.convert %30721 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1752)
      %33233 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1753)
      %33234 = stablehlo.convert %33233 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1754)
      %33235 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33236 = stablehlo.broadcast_in_dim %33235, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1755)
      %33237 = stablehlo.subtract %33236, %33234 : tensor<1x1x2048x1xf32> loc(#loc1755)
      %33238 = stablehlo.convert %33237 : tensor<1x1x2048x1xf32> loc(#loc1756)
      %33239 = stablehlo.broadcast_in_dim %33238, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33240 = stablehlo.multiply %33232, %33239 : tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33241 = stablehlo.abs %33240 : tensor<1x1x2048x2048xf32> loc(#loc1758)
      %33242 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1759)
      %33243 = stablehlo.reduce(%33241 init: %33242) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1759)
        stablehlo.return %33956 : tensor<f32> loc(#loc1759)
      } loc(#loc1759)
      %33244 = func.call @_where_437(%30042, %33243, %iterArg_464) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33245 = stablehlo.convert %33237 : tensor<1x1x2048x1xf32> loc(#loc1756)
      %33246 = stablehlo.broadcast_in_dim %33245, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33247 = stablehlo.multiply %33232, %33246 : tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33248 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1760)
      %33249 = stablehlo.reduce(%33247 init: %33248) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1760)
        stablehlo.return %33956 : tensor<f32> loc(#loc1760)
      } loc(#loc1760)
      %33250 = stablehlo.convert %33237 : tensor<1x1x2048x1xf32> loc(#loc1756)
      %33251 = stablehlo.broadcast_in_dim %33250, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33252 = stablehlo.multiply %iterArg_465, %33251 : tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33253 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1760)
      %33254 = stablehlo.reduce(%33252 init: %33253) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1760)
        stablehlo.return %33956 : tensor<f32> loc(#loc1760)
      } loc(#loc1760)
      %33255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33256 = stablehlo.broadcast_in_dim %33255, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1761)
      %33257 = stablehlo.maximum %33256, %33254 : tensor<1xf32> loc(#loc1761)
      %33258 = stablehlo.divide %33249, %33257 : tensor<1xf32> loc(#loc1762)
      %33259 = func.call @_where_438(%30042, %33258, %iterArg_466) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33260 = stablehlo.broadcast_in_dim %33258, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1763)
      %33261 = stablehlo.broadcast_in_dim %33260, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1764)
      %33262 = stablehlo.subtract %33232, %33261 : tensor<1x1x2048x2048xf32> loc(#loc1764)
      %33263 = stablehlo.convert %33237 : tensor<1x1x2048x1xf32> loc(#loc1756)
      %33264 = stablehlo.broadcast_in_dim %33263, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33265 = stablehlo.multiply %33262, %33264 : tensor<1x1x2048x2048xf32> loc(#loc1757)
      %33266 = stablehlo.multiply %33265, %33265 : tensor<1x1x2048x2048xf32> loc(#loc1765)
      %33267 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1760)
      %33268 = stablehlo.reduce(%33266 init: %33267) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1760)
        stablehlo.return %33956 : tensor<f32> loc(#loc1760)
      } loc(#loc1760)
      %33269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33270 = stablehlo.broadcast_in_dim %33269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1761)
      %33271 = stablehlo.maximum %33270, %33254 : tensor<1xf32> loc(#loc1761)
      %33272 = stablehlo.divide %33268, %33271 : tensor<1xf32> loc(#loc1762)
      %33273 = stablehlo.sqrt %33272 : tensor<1xf32> loc(#loc1766)
      %33274 = func.call @_where_439(%30042, %33273, %iterArg_467) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33275 = stablehlo.convert %30894 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1767)
      %33276 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1768)
      %33277 = stablehlo.convert %33276 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1769)
      %33278 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33279 = stablehlo.broadcast_in_dim %33278, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1770)
      %33280 = stablehlo.subtract %33279, %33277 : tensor<1x1x2048x1xf32> loc(#loc1770)
      %33281 = stablehlo.convert %33280 : tensor<1x1x2048x1xf32> loc(#loc1771)
      %33282 = stablehlo.broadcast_in_dim %33281, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33283 = stablehlo.multiply %33275, %33282 : tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33284 = stablehlo.abs %33283 : tensor<1x1x2048x2048xf32> loc(#loc1773)
      %33285 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1774)
      %33286 = stablehlo.reduce(%33284 init: %33285) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1774)
        stablehlo.return %33956 : tensor<f32> loc(#loc1774)
      } loc(#loc1774)
      %33287 = func.call @_where_440(%30042, %33286, %iterArg_468) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33288 = stablehlo.convert %33280 : tensor<1x1x2048x1xf32> loc(#loc1771)
      %33289 = stablehlo.broadcast_in_dim %33288, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33290 = stablehlo.multiply %33275, %33289 : tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33291 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1775)
      %33292 = stablehlo.reduce(%33290 init: %33291) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1775)
        stablehlo.return %33956 : tensor<f32> loc(#loc1775)
      } loc(#loc1775)
      %33293 = stablehlo.convert %33280 : tensor<1x1x2048x1xf32> loc(#loc1771)
      %33294 = stablehlo.broadcast_in_dim %33293, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33295 = stablehlo.multiply %iterArg_469, %33294 : tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33296 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1775)
      %33297 = stablehlo.reduce(%33295 init: %33296) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1775)
        stablehlo.return %33956 : tensor<f32> loc(#loc1775)
      } loc(#loc1775)
      %33298 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33299 = stablehlo.broadcast_in_dim %33298, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1776)
      %33300 = stablehlo.maximum %33299, %33297 : tensor<1xf32> loc(#loc1776)
      %33301 = stablehlo.divide %33292, %33300 : tensor<1xf32> loc(#loc1777)
      %33302 = func.call @_where_441(%30042, %33301, %iterArg_470) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33303 = stablehlo.broadcast_in_dim %33301, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1778)
      %33304 = stablehlo.broadcast_in_dim %33303, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1779)
      %33305 = stablehlo.subtract %33275, %33304 : tensor<1x1x2048x2048xf32> loc(#loc1779)
      %33306 = stablehlo.convert %33280 : tensor<1x1x2048x1xf32> loc(#loc1771)
      %33307 = stablehlo.broadcast_in_dim %33306, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33308 = stablehlo.multiply %33305, %33307 : tensor<1x1x2048x2048xf32> loc(#loc1772)
      %33309 = stablehlo.multiply %33308, %33308 : tensor<1x1x2048x2048xf32> loc(#loc1780)
      %33310 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1775)
      %33311 = stablehlo.reduce(%33309 init: %33310) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1775)
        stablehlo.return %33956 : tensor<f32> loc(#loc1775)
      } loc(#loc1775)
      %33312 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33313 = stablehlo.broadcast_in_dim %33312, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1776)
      %33314 = stablehlo.maximum %33313, %33297 : tensor<1xf32> loc(#loc1776)
      %33315 = stablehlo.divide %33311, %33314 : tensor<1xf32> loc(#loc1777)
      %33316 = stablehlo.sqrt %33315 : tensor<1xf32> loc(#loc1781)
      %33317 = func.call @_where_442(%30042, %33316, %iterArg_471) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33318 = stablehlo.convert %31067 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1782)
      %33319 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1783)
      %33320 = stablehlo.convert %33319 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1784)
      %33321 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33322 = stablehlo.broadcast_in_dim %33321, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1785)
      %33323 = stablehlo.subtract %33322, %33320 : tensor<1x1x2048x1xf32> loc(#loc1785)
      %33324 = stablehlo.convert %33323 : tensor<1x1x2048x1xf32> loc(#loc1786)
      %33325 = stablehlo.broadcast_in_dim %33324, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33326 = stablehlo.multiply %33318, %33325 : tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33327 = stablehlo.abs %33326 : tensor<1x1x2048x2048xf32> loc(#loc1788)
      %33328 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1789)
      %33329 = stablehlo.reduce(%33327 init: %33328) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1789)
        stablehlo.return %33956 : tensor<f32> loc(#loc1789)
      } loc(#loc1789)
      %33330 = func.call @_where_443(%30042, %33329, %iterArg_472) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33331 = stablehlo.convert %33323 : tensor<1x1x2048x1xf32> loc(#loc1786)
      %33332 = stablehlo.broadcast_in_dim %33331, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33333 = stablehlo.multiply %33318, %33332 : tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33334 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1790)
      %33335 = stablehlo.reduce(%33333 init: %33334) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1790)
        stablehlo.return %33956 : tensor<f32> loc(#loc1790)
      } loc(#loc1790)
      %33336 = stablehlo.convert %33323 : tensor<1x1x2048x1xf32> loc(#loc1786)
      %33337 = stablehlo.broadcast_in_dim %33336, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33338 = stablehlo.multiply %iterArg_473, %33337 : tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33339 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1790)
      %33340 = stablehlo.reduce(%33338 init: %33339) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1790)
        stablehlo.return %33956 : tensor<f32> loc(#loc1790)
      } loc(#loc1790)
      %33341 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33342 = stablehlo.broadcast_in_dim %33341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1791)
      %33343 = stablehlo.maximum %33342, %33340 : tensor<1xf32> loc(#loc1791)
      %33344 = stablehlo.divide %33335, %33343 : tensor<1xf32> loc(#loc1792)
      %33345 = func.call @_where_444(%30042, %33344, %iterArg_474) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33346 = stablehlo.broadcast_in_dim %33344, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1793)
      %33347 = stablehlo.broadcast_in_dim %33346, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1794)
      %33348 = stablehlo.subtract %33318, %33347 : tensor<1x1x2048x2048xf32> loc(#loc1794)
      %33349 = stablehlo.convert %33323 : tensor<1x1x2048x1xf32> loc(#loc1786)
      %33350 = stablehlo.broadcast_in_dim %33349, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33351 = stablehlo.multiply %33348, %33350 : tensor<1x1x2048x2048xf32> loc(#loc1787)
      %33352 = stablehlo.multiply %33351, %33351 : tensor<1x1x2048x2048xf32> loc(#loc1795)
      %33353 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1790)
      %33354 = stablehlo.reduce(%33352 init: %33353) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1790)
        stablehlo.return %33956 : tensor<f32> loc(#loc1790)
      } loc(#loc1790)
      %33355 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33356 = stablehlo.broadcast_in_dim %33355, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1791)
      %33357 = stablehlo.maximum %33356, %33340 : tensor<1xf32> loc(#loc1791)
      %33358 = stablehlo.divide %33354, %33357 : tensor<1xf32> loc(#loc1792)
      %33359 = stablehlo.sqrt %33358 : tensor<1xf32> loc(#loc1796)
      %33360 = func.call @_where_445(%30042, %33359, %iterArg_475) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33361 = stablehlo.convert %31240 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1797)
      %33362 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1798)
      %33363 = stablehlo.convert %33362 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1799)
      %33364 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33365 = stablehlo.broadcast_in_dim %33364, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1800)
      %33366 = stablehlo.subtract %33365, %33363 : tensor<1x1x2048x1xf32> loc(#loc1800)
      %33367 = stablehlo.convert %33366 : tensor<1x1x2048x1xf32> loc(#loc1801)
      %33368 = stablehlo.broadcast_in_dim %33367, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33369 = stablehlo.multiply %33361, %33368 : tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33370 = stablehlo.abs %33369 : tensor<1x1x2048x2048xf32> loc(#loc1803)
      %33371 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1804)
      %33372 = stablehlo.reduce(%33370 init: %33371) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1804)
        stablehlo.return %33956 : tensor<f32> loc(#loc1804)
      } loc(#loc1804)
      %33373 = func.call @_where_446(%30042, %33372, %iterArg_476) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33374 = stablehlo.convert %33366 : tensor<1x1x2048x1xf32> loc(#loc1801)
      %33375 = stablehlo.broadcast_in_dim %33374, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33376 = stablehlo.multiply %33361, %33375 : tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33377 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1805)
      %33378 = stablehlo.reduce(%33376 init: %33377) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1805)
        stablehlo.return %33956 : tensor<f32> loc(#loc1805)
      } loc(#loc1805)
      %33379 = stablehlo.convert %33366 : tensor<1x1x2048x1xf32> loc(#loc1801)
      %33380 = stablehlo.broadcast_in_dim %33379, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33381 = stablehlo.multiply %iterArg_477, %33380 : tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33382 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1805)
      %33383 = stablehlo.reduce(%33381 init: %33382) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1805)
        stablehlo.return %33956 : tensor<f32> loc(#loc1805)
      } loc(#loc1805)
      %33384 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33385 = stablehlo.broadcast_in_dim %33384, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1806)
      %33386 = stablehlo.maximum %33385, %33383 : tensor<1xf32> loc(#loc1806)
      %33387 = stablehlo.divide %33378, %33386 : tensor<1xf32> loc(#loc1807)
      %33388 = func.call @_where_447(%30042, %33387, %iterArg_478) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33389 = stablehlo.broadcast_in_dim %33387, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1808)
      %33390 = stablehlo.broadcast_in_dim %33389, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1809)
      %33391 = stablehlo.subtract %33361, %33390 : tensor<1x1x2048x2048xf32> loc(#loc1809)
      %33392 = stablehlo.convert %33366 : tensor<1x1x2048x1xf32> loc(#loc1801)
      %33393 = stablehlo.broadcast_in_dim %33392, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33394 = stablehlo.multiply %33391, %33393 : tensor<1x1x2048x2048xf32> loc(#loc1802)
      %33395 = stablehlo.multiply %33394, %33394 : tensor<1x1x2048x2048xf32> loc(#loc1810)
      %33396 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1805)
      %33397 = stablehlo.reduce(%33395 init: %33396) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1805)
        stablehlo.return %33956 : tensor<f32> loc(#loc1805)
      } loc(#loc1805)
      %33398 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33399 = stablehlo.broadcast_in_dim %33398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1806)
      %33400 = stablehlo.maximum %33399, %33383 : tensor<1xf32> loc(#loc1806)
      %33401 = stablehlo.divide %33397, %33400 : tensor<1xf32> loc(#loc1807)
      %33402 = stablehlo.sqrt %33401 : tensor<1xf32> loc(#loc1811)
      %33403 = func.call @_where_448(%30042, %33402, %iterArg_479) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33404 = stablehlo.convert %31413 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1812)
      %33405 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1813)
      %33406 = stablehlo.convert %33405 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1814)
      %33407 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33408 = stablehlo.broadcast_in_dim %33407, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1815)
      %33409 = stablehlo.subtract %33408, %33406 : tensor<1x1x2048x1xf32> loc(#loc1815)
      %33410 = stablehlo.convert %33409 : tensor<1x1x2048x1xf32> loc(#loc1816)
      %33411 = stablehlo.broadcast_in_dim %33410, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33412 = stablehlo.multiply %33404, %33411 : tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33413 = stablehlo.abs %33412 : tensor<1x1x2048x2048xf32> loc(#loc1818)
      %33414 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1819)
      %33415 = stablehlo.reduce(%33413 init: %33414) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1819)
        stablehlo.return %33956 : tensor<f32> loc(#loc1819)
      } loc(#loc1819)
      %33416 = func.call @_where_449(%30042, %33415, %iterArg_480) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33417 = stablehlo.convert %33409 : tensor<1x1x2048x1xf32> loc(#loc1816)
      %33418 = stablehlo.broadcast_in_dim %33417, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33419 = stablehlo.multiply %33404, %33418 : tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33420 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1820)
      %33421 = stablehlo.reduce(%33419 init: %33420) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1820)
        stablehlo.return %33956 : tensor<f32> loc(#loc1820)
      } loc(#loc1820)
      %33422 = stablehlo.convert %33409 : tensor<1x1x2048x1xf32> loc(#loc1816)
      %33423 = stablehlo.broadcast_in_dim %33422, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33424 = stablehlo.multiply %iterArg_481, %33423 : tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33425 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1820)
      %33426 = stablehlo.reduce(%33424 init: %33425) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1820)
        stablehlo.return %33956 : tensor<f32> loc(#loc1820)
      } loc(#loc1820)
      %33427 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33428 = stablehlo.broadcast_in_dim %33427, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1821)
      %33429 = stablehlo.maximum %33428, %33426 : tensor<1xf32> loc(#loc1821)
      %33430 = stablehlo.divide %33421, %33429 : tensor<1xf32> loc(#loc1822)
      %33431 = func.call @_where_450(%30042, %33430, %iterArg_482) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33432 = stablehlo.broadcast_in_dim %33430, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1823)
      %33433 = stablehlo.broadcast_in_dim %33432, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1824)
      %33434 = stablehlo.subtract %33404, %33433 : tensor<1x1x2048x2048xf32> loc(#loc1824)
      %33435 = stablehlo.convert %33409 : tensor<1x1x2048x1xf32> loc(#loc1816)
      %33436 = stablehlo.broadcast_in_dim %33435, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33437 = stablehlo.multiply %33434, %33436 : tensor<1x1x2048x2048xf32> loc(#loc1817)
      %33438 = stablehlo.multiply %33437, %33437 : tensor<1x1x2048x2048xf32> loc(#loc1825)
      %33439 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1820)
      %33440 = stablehlo.reduce(%33438 init: %33439) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1820)
        stablehlo.return %33956 : tensor<f32> loc(#loc1820)
      } loc(#loc1820)
      %33441 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33442 = stablehlo.broadcast_in_dim %33441, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1821)
      %33443 = stablehlo.maximum %33442, %33426 : tensor<1xf32> loc(#loc1821)
      %33444 = stablehlo.divide %33440, %33443 : tensor<1xf32> loc(#loc1822)
      %33445 = stablehlo.sqrt %33444 : tensor<1xf32> loc(#loc1826)
      %33446 = func.call @_where_451(%30042, %33445, %iterArg_483) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33447 = stablehlo.convert %31586 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc1827)
      %33448 = stablehlo.broadcast_in_dim %30176, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc1828)
      %33449 = stablehlo.convert %33448 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc1829)
      %33450 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33451 = stablehlo.broadcast_in_dim %33450, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc1830)
      %33452 = stablehlo.subtract %33451, %33449 : tensor<1x1x2048x1xf32> loc(#loc1830)
      %33453 = stablehlo.convert %33452 : tensor<1x1x2048x1xf32> loc(#loc1831)
      %33454 = stablehlo.broadcast_in_dim %33453, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33455 = stablehlo.multiply %33447, %33454 : tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33456 = stablehlo.abs %33455 : tensor<1x1x2048x2048xf32> loc(#loc1833)
      %33457 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1834)
      %33458 = stablehlo.reduce(%33456 init: %33457) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc1834)
        stablehlo.return %33956 : tensor<f32> loc(#loc1834)
      } loc(#loc1834)
      %33459 = func.call @_where_452(%30042, %33458, %iterArg_484) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33460 = stablehlo.convert %33452 : tensor<1x1x2048x1xf32> loc(#loc1831)
      %33461 = stablehlo.broadcast_in_dim %33460, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33462 = stablehlo.multiply %33447, %33461 : tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33463 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1835)
      %33464 = stablehlo.reduce(%33462 init: %33463) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1835)
        stablehlo.return %33956 : tensor<f32> loc(#loc1835)
      } loc(#loc1835)
      %33465 = stablehlo.convert %33452 : tensor<1x1x2048x1xf32> loc(#loc1831)
      %33466 = stablehlo.broadcast_in_dim %33465, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33467 = stablehlo.multiply %iterArg_485, %33466 : tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33468 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1835)
      %33469 = stablehlo.reduce(%33467 init: %33468) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1835)
        stablehlo.return %33956 : tensor<f32> loc(#loc1835)
      } loc(#loc1835)
      %33470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33471 = stablehlo.broadcast_in_dim %33470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1836)
      %33472 = stablehlo.maximum %33471, %33469 : tensor<1xf32> loc(#loc1836)
      %33473 = stablehlo.divide %33464, %33472 : tensor<1xf32> loc(#loc1837)
      %33474 = func.call @_where_453(%30042, %33473, %iterArg_486) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33475 = stablehlo.broadcast_in_dim %33473, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc1838)
      %33476 = stablehlo.broadcast_in_dim %33475, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1839)
      %33477 = stablehlo.subtract %33447, %33476 : tensor<1x1x2048x2048xf32> loc(#loc1839)
      %33478 = stablehlo.convert %33452 : tensor<1x1x2048x1xf32> loc(#loc1831)
      %33479 = stablehlo.broadcast_in_dim %33478, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33480 = stablehlo.multiply %33477, %33479 : tensor<1x1x2048x2048xf32> loc(#loc1832)
      %33481 = stablehlo.multiply %33480, %33480 : tensor<1x1x2048x2048xf32> loc(#loc1840)
      %33482 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1835)
      %33483 = stablehlo.reduce(%33481 init: %33482) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %33956 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1835)
        stablehlo.return %33956 : tensor<f32> loc(#loc1835)
      } loc(#loc1835)
      %33484 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc302)
      %33485 = stablehlo.broadcast_in_dim %33484, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1836)
      %33486 = stablehlo.maximum %33485, %33469 : tensor<1xf32> loc(#loc1836)
      %33487 = stablehlo.divide %33483, %33486 : tensor<1xf32> loc(#loc1837)
      %33488 = stablehlo.sqrt %33487 : tensor<1xf32> loc(#loc1841)
      %33489 = func.call @_where_454(%30042, %33488, %iterArg_487) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc1609)
      %33490 = stablehlo.broadcast_in_dim %32797, dims = [1, 2, 3, 4] : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1842)
      %33491 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33492 = stablehlo.compare  LT, %iterArg_488, %33491,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33493 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33494 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33495 = stablehlo.add %33493, %33494 : tensor<i32> loc(#loc1845)
      %33496 = stablehlo.select %33492, %33495, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33497 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33498 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33499 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33500 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33501 = stablehlo.dynamic_update_slice %iterArg_491, %33490, %33496, %33497, %33498, %33499, %33500 : (tensor<1x1x1x2048x2048xbf16>, tensor<1x1x1x2048x2048xbf16>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1847)
      %33502 = stablehlo.broadcast_in_dim %32814, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33503 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33504 = stablehlo.compare  LT, %iterArg_488, %33503,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33505 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33506 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33507 = stablehlo.add %33505, %33506 : tensor<i32> loc(#loc1845)
      %33508 = stablehlo.select %33504, %33507, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33509 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33510 = stablehlo.dynamic_update_slice %iterArg_492, %33502, %33508, %33509 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33511 = stablehlo.broadcast_in_dim %32829, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33512 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33513 = stablehlo.compare  LT, %iterArg_488, %33512,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33514 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33515 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33516 = stablehlo.add %33514, %33515 : tensor<i32> loc(#loc1845)
      %33517 = stablehlo.select %33513, %33516, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33518 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33519 = stablehlo.dynamic_update_slice %iterArg_493, %33511, %33517, %33518 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33520 = stablehlo.broadcast_in_dim %32844, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33521 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33522 = stablehlo.compare  LT, %iterArg_488, %33521,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33523 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33524 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33525 = stablehlo.add %33523, %33524 : tensor<i32> loc(#loc1845)
      %33526 = stablehlo.select %33522, %33525, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33527 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33528 = stablehlo.dynamic_update_slice %iterArg_494, %33520, %33526, %33527 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33529 = stablehlo.broadcast_in_dim %32857, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33530 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33531 = stablehlo.compare  LT, %iterArg_488, %33530,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33532 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33533 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33534 = stablehlo.add %33532, %33533 : tensor<i32> loc(#loc1845)
      %33535 = stablehlo.select %33531, %33534, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33536 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33537 = stablehlo.dynamic_update_slice %iterArg_495, %33529, %33535, %33536 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33538 = stablehlo.broadcast_in_dim %32872, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33539 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33540 = stablehlo.compare  LT, %iterArg_488, %33539,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33541 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33542 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33543 = stablehlo.add %33541, %33542 : tensor<i32> loc(#loc1845)
      %33544 = stablehlo.select %33540, %33543, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33545 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33546 = stablehlo.dynamic_update_slice %iterArg_496, %33538, %33544, %33545 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33547 = stablehlo.broadcast_in_dim %32887, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33548 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33549 = stablehlo.compare  LT, %iterArg_488, %33548,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33550 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33551 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33552 = stablehlo.add %33550, %33551 : tensor<i32> loc(#loc1845)
      %33553 = stablehlo.select %33549, %33552, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33554 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33555 = stablehlo.dynamic_update_slice %iterArg_497, %33547, %33553, %33554 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33556 = stablehlo.broadcast_in_dim %32900, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33557 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33558 = stablehlo.compare  LT, %iterArg_488, %33557,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33559 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33560 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33561 = stablehlo.add %33559, %33560 : tensor<i32> loc(#loc1845)
      %33562 = stablehlo.select %33558, %33561, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33563 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33564 = stablehlo.dynamic_update_slice %iterArg_498, %33556, %33562, %33563 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33565 = stablehlo.broadcast_in_dim %32915, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33566 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33567 = stablehlo.compare  LT, %iterArg_488, %33566,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33568 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33569 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33570 = stablehlo.add %33568, %33569 : tensor<i32> loc(#loc1845)
      %33571 = stablehlo.select %33567, %33570, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33572 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33573 = stablehlo.dynamic_update_slice %iterArg_499, %33565, %33571, %33572 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33574 = stablehlo.broadcast_in_dim %32930, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33575 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33576 = stablehlo.compare  LT, %iterArg_488, %33575,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33577 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33578 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33579 = stablehlo.add %33577, %33578 : tensor<i32> loc(#loc1845)
      %33580 = stablehlo.select %33576, %33579, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33581 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33582 = stablehlo.dynamic_update_slice %iterArg_500, %33574, %33580, %33581 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33583 = stablehlo.broadcast_in_dim %32943, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33584 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33585 = stablehlo.compare  LT, %iterArg_488, %33584,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33586 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33587 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33588 = stablehlo.add %33586, %33587 : tensor<i32> loc(#loc1845)
      %33589 = stablehlo.select %33585, %33588, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33590 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33591 = stablehlo.dynamic_update_slice %iterArg_501, %33583, %33589, %33590 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33592 = stablehlo.broadcast_in_dim %32958, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33593 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33594 = stablehlo.compare  LT, %iterArg_488, %33593,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33595 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33596 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33597 = stablehlo.add %33595, %33596 : tensor<i32> loc(#loc1845)
      %33598 = stablehlo.select %33594, %33597, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33599 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33600 = stablehlo.dynamic_update_slice %iterArg_502, %33592, %33598, %33599 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33601 = stablehlo.broadcast_in_dim %32973, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33602 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33603 = stablehlo.compare  LT, %iterArg_488, %33602,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33604 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33605 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33606 = stablehlo.add %33604, %33605 : tensor<i32> loc(#loc1845)
      %33607 = stablehlo.select %33603, %33606, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33608 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33609 = stablehlo.dynamic_update_slice %iterArg_503, %33601, %33607, %33608 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33610 = stablehlo.broadcast_in_dim %32986, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33611 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33612 = stablehlo.compare  LT, %iterArg_488, %33611,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33613 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33614 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33615 = stablehlo.add %33613, %33614 : tensor<i32> loc(#loc1845)
      %33616 = stablehlo.select %33612, %33615, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33617 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33618 = stablehlo.dynamic_update_slice %iterArg_504, %33610, %33616, %33617 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33619 = stablehlo.broadcast_in_dim %33001, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33620 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33621 = stablehlo.compare  LT, %iterArg_488, %33620,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33622 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33623 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33624 = stablehlo.add %33622, %33623 : tensor<i32> loc(#loc1845)
      %33625 = stablehlo.select %33621, %33624, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33626 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33627 = stablehlo.dynamic_update_slice %iterArg_505, %33619, %33625, %33626 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33628 = stablehlo.broadcast_in_dim %33016, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33629 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33630 = stablehlo.compare  LT, %iterArg_488, %33629,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33631 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33632 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33633 = stablehlo.add %33631, %33632 : tensor<i32> loc(#loc1845)
      %33634 = stablehlo.select %33630, %33633, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33635 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33636 = stablehlo.dynamic_update_slice %iterArg_506, %33628, %33634, %33635 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33637 = stablehlo.broadcast_in_dim %33029, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33638 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33639 = stablehlo.compare  LT, %iterArg_488, %33638,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33640 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33641 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33642 = stablehlo.add %33640, %33641 : tensor<i32> loc(#loc1845)
      %33643 = stablehlo.select %33639, %33642, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33644 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33645 = stablehlo.dynamic_update_slice %iterArg_507, %33637, %33643, %33644 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33646 = stablehlo.broadcast_in_dim %33044, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33647 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33648 = stablehlo.compare  LT, %iterArg_488, %33647,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33649 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33650 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33651 = stablehlo.add %33649, %33650 : tensor<i32> loc(#loc1845)
      %33652 = stablehlo.select %33648, %33651, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33653 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33654 = stablehlo.dynamic_update_slice %iterArg_508, %33646, %33652, %33653 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33655 = stablehlo.broadcast_in_dim %33059, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33656 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33657 = stablehlo.compare  LT, %iterArg_488, %33656,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33658 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33659 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33660 = stablehlo.add %33658, %33659 : tensor<i32> loc(#loc1845)
      %33661 = stablehlo.select %33657, %33660, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33662 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33663 = stablehlo.dynamic_update_slice %iterArg_509, %33655, %33661, %33662 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33664 = stablehlo.broadcast_in_dim %33072, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33665 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33666 = stablehlo.compare  LT, %iterArg_488, %33665,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33667 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33668 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33669 = stablehlo.add %33667, %33668 : tensor<i32> loc(#loc1845)
      %33670 = stablehlo.select %33666, %33669, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33671 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33672 = stablehlo.dynamic_update_slice %iterArg_510, %33664, %33670, %33671 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33673 = stablehlo.broadcast_in_dim %33087, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33674 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33675 = stablehlo.compare  LT, %iterArg_488, %33674,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33676 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33677 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33678 = stablehlo.add %33676, %33677 : tensor<i32> loc(#loc1845)
      %33679 = stablehlo.select %33675, %33678, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33680 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33681 = stablehlo.dynamic_update_slice %iterArg_511, %33673, %33679, %33680 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33682 = stablehlo.broadcast_in_dim %33102, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33683 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33684 = stablehlo.compare  LT, %iterArg_488, %33683,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33685 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33686 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33687 = stablehlo.add %33685, %33686 : tensor<i32> loc(#loc1845)
      %33688 = stablehlo.select %33684, %33687, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33689 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33690 = stablehlo.dynamic_update_slice %iterArg_512, %33682, %33688, %33689 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33691 = stablehlo.broadcast_in_dim %33115, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33692 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33693 = stablehlo.compare  LT, %iterArg_488, %33692,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33694 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33695 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33696 = stablehlo.add %33694, %33695 : tensor<i32> loc(#loc1845)
      %33697 = stablehlo.select %33693, %33696, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33698 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33699 = stablehlo.dynamic_update_slice %iterArg_513, %33691, %33697, %33698 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33700 = stablehlo.broadcast_in_dim %33130, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33701 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33702 = stablehlo.compare  LT, %iterArg_488, %33701,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33703 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33704 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33705 = stablehlo.add %33703, %33704 : tensor<i32> loc(#loc1845)
      %33706 = stablehlo.select %33702, %33705, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33707 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33708 = stablehlo.dynamic_update_slice %iterArg_514, %33700, %33706, %33707 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33709 = stablehlo.broadcast_in_dim %33145, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33710 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33711 = stablehlo.compare  LT, %iterArg_488, %33710,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33712 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33713 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33714 = stablehlo.add %33712, %33713 : tensor<i32> loc(#loc1845)
      %33715 = stablehlo.select %33711, %33714, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33716 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33717 = stablehlo.dynamic_update_slice %iterArg_515, %33709, %33715, %33716 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33718 = stablehlo.broadcast_in_dim %33158, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33719 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33720 = stablehlo.compare  LT, %iterArg_488, %33719,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33721 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33722 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33723 = stablehlo.add %33721, %33722 : tensor<i32> loc(#loc1845)
      %33724 = stablehlo.select %33720, %33723, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33725 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33726 = stablehlo.dynamic_update_slice %iterArg_516, %33718, %33724, %33725 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33727 = stablehlo.broadcast_in_dim %33173, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33728 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33729 = stablehlo.compare  LT, %iterArg_488, %33728,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33730 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33731 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33732 = stablehlo.add %33730, %33731 : tensor<i32> loc(#loc1845)
      %33733 = stablehlo.select %33729, %33732, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33734 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33735 = stablehlo.dynamic_update_slice %iterArg_517, %33727, %33733, %33734 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33736 = stablehlo.broadcast_in_dim %33188, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33737 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33738 = stablehlo.compare  LT, %iterArg_488, %33737,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33739 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33740 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33741 = stablehlo.add %33739, %33740 : tensor<i32> loc(#loc1845)
      %33742 = stablehlo.select %33738, %33741, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33743 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33744 = stablehlo.dynamic_update_slice %iterArg_518, %33736, %33742, %33743 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33745 = stablehlo.broadcast_in_dim %33201, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33746 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33747 = stablehlo.compare  LT, %iterArg_488, %33746,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33748 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33749 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33750 = stablehlo.add %33748, %33749 : tensor<i32> loc(#loc1845)
      %33751 = stablehlo.select %33747, %33750, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33752 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33753 = stablehlo.dynamic_update_slice %iterArg_519, %33745, %33751, %33752 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33754 = stablehlo.broadcast_in_dim %33216, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33755 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33756 = stablehlo.compare  LT, %iterArg_488, %33755,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33757 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33758 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33759 = stablehlo.add %33757, %33758 : tensor<i32> loc(#loc1845)
      %33760 = stablehlo.select %33756, %33759, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33761 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33762 = stablehlo.dynamic_update_slice %iterArg_520, %33754, %33760, %33761 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33763 = stablehlo.broadcast_in_dim %33231, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33764 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33765 = stablehlo.compare  LT, %iterArg_488, %33764,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33766 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33767 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33768 = stablehlo.add %33766, %33767 : tensor<i32> loc(#loc1845)
      %33769 = stablehlo.select %33765, %33768, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33770 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33771 = stablehlo.dynamic_update_slice %iterArg_521, %33763, %33769, %33770 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33772 = stablehlo.broadcast_in_dim %33244, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33773 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33774 = stablehlo.compare  LT, %iterArg_488, %33773,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33775 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33776 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33777 = stablehlo.add %33775, %33776 : tensor<i32> loc(#loc1845)
      %33778 = stablehlo.select %33774, %33777, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33779 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33780 = stablehlo.dynamic_update_slice %iterArg_522, %33772, %33778, %33779 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33781 = stablehlo.broadcast_in_dim %33259, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33782 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33783 = stablehlo.compare  LT, %iterArg_488, %33782,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33784 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33785 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33786 = stablehlo.add %33784, %33785 : tensor<i32> loc(#loc1845)
      %33787 = stablehlo.select %33783, %33786, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33788 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33789 = stablehlo.dynamic_update_slice %iterArg_523, %33781, %33787, %33788 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33790 = stablehlo.broadcast_in_dim %33274, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33791 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33792 = stablehlo.compare  LT, %iterArg_488, %33791,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33793 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33794 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33795 = stablehlo.add %33793, %33794 : tensor<i32> loc(#loc1845)
      %33796 = stablehlo.select %33792, %33795, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33797 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33798 = stablehlo.dynamic_update_slice %iterArg_524, %33790, %33796, %33797 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33799 = stablehlo.broadcast_in_dim %33287, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33800 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33801 = stablehlo.compare  LT, %iterArg_488, %33800,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33802 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33803 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33804 = stablehlo.add %33802, %33803 : tensor<i32> loc(#loc1845)
      %33805 = stablehlo.select %33801, %33804, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33806 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33807 = stablehlo.dynamic_update_slice %iterArg_525, %33799, %33805, %33806 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33808 = stablehlo.broadcast_in_dim %33302, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33809 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33810 = stablehlo.compare  LT, %iterArg_488, %33809,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33811 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33812 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33813 = stablehlo.add %33811, %33812 : tensor<i32> loc(#loc1845)
      %33814 = stablehlo.select %33810, %33813, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33815 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33816 = stablehlo.dynamic_update_slice %iterArg_526, %33808, %33814, %33815 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33817 = stablehlo.broadcast_in_dim %33317, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33818 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33819 = stablehlo.compare  LT, %iterArg_488, %33818,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33820 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33821 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33822 = stablehlo.add %33820, %33821 : tensor<i32> loc(#loc1845)
      %33823 = stablehlo.select %33819, %33822, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33824 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33825 = stablehlo.dynamic_update_slice %iterArg_527, %33817, %33823, %33824 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33826 = stablehlo.broadcast_in_dim %33330, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33827 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33828 = stablehlo.compare  LT, %iterArg_488, %33827,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33829 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33830 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33831 = stablehlo.add %33829, %33830 : tensor<i32> loc(#loc1845)
      %33832 = stablehlo.select %33828, %33831, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33833 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33834 = stablehlo.dynamic_update_slice %iterArg_528, %33826, %33832, %33833 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33835 = stablehlo.broadcast_in_dim %33345, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33836 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33837 = stablehlo.compare  LT, %iterArg_488, %33836,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33838 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33839 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33840 = stablehlo.add %33838, %33839 : tensor<i32> loc(#loc1845)
      %33841 = stablehlo.select %33837, %33840, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33842 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33843 = stablehlo.dynamic_update_slice %iterArg_529, %33835, %33841, %33842 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33844 = stablehlo.broadcast_in_dim %33360, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33845 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33846 = stablehlo.compare  LT, %iterArg_488, %33845,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33847 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33848 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33849 = stablehlo.add %33847, %33848 : tensor<i32> loc(#loc1845)
      %33850 = stablehlo.select %33846, %33849, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33851 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33852 = stablehlo.dynamic_update_slice %iterArg_530, %33844, %33850, %33851 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33853 = stablehlo.broadcast_in_dim %33373, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33854 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33855 = stablehlo.compare  LT, %iterArg_488, %33854,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33856 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33857 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33858 = stablehlo.add %33856, %33857 : tensor<i32> loc(#loc1845)
      %33859 = stablehlo.select %33855, %33858, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33860 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33861 = stablehlo.dynamic_update_slice %iterArg_531, %33853, %33859, %33860 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33862 = stablehlo.broadcast_in_dim %33388, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33863 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33864 = stablehlo.compare  LT, %iterArg_488, %33863,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33865 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33866 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33867 = stablehlo.add %33865, %33866 : tensor<i32> loc(#loc1845)
      %33868 = stablehlo.select %33864, %33867, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33869 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33870 = stablehlo.dynamic_update_slice %iterArg_532, %33862, %33868, %33869 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33871 = stablehlo.broadcast_in_dim %33403, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33872 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33873 = stablehlo.compare  LT, %iterArg_488, %33872,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33874 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33875 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33876 = stablehlo.add %33874, %33875 : tensor<i32> loc(#loc1845)
      %33877 = stablehlo.select %33873, %33876, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33878 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33879 = stablehlo.dynamic_update_slice %iterArg_533, %33871, %33877, %33878 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33880 = stablehlo.broadcast_in_dim %33416, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33881 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33882 = stablehlo.compare  LT, %iterArg_488, %33881,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33883 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33884 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33885 = stablehlo.add %33883, %33884 : tensor<i32> loc(#loc1845)
      %33886 = stablehlo.select %33882, %33885, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33887 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33888 = stablehlo.dynamic_update_slice %iterArg_534, %33880, %33886, %33887 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33889 = stablehlo.broadcast_in_dim %33431, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33890 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33891 = stablehlo.compare  LT, %iterArg_488, %33890,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33892 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33893 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33894 = stablehlo.add %33892, %33893 : tensor<i32> loc(#loc1845)
      %33895 = stablehlo.select %33891, %33894, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33896 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33897 = stablehlo.dynamic_update_slice %iterArg_535, %33889, %33895, %33896 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33898 = stablehlo.broadcast_in_dim %33446, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33899 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33900 = stablehlo.compare  LT, %iterArg_488, %33899,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33901 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33902 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33903 = stablehlo.add %33901, %33902 : tensor<i32> loc(#loc1845)
      %33904 = stablehlo.select %33900, %33903, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33905 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33906 = stablehlo.dynamic_update_slice %iterArg_536, %33898, %33904, %33905 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33907 = stablehlo.broadcast_in_dim %33459, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33908 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33909 = stablehlo.compare  LT, %iterArg_488, %33908,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33910 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33911 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33912 = stablehlo.add %33910, %33911 : tensor<i32> loc(#loc1845)
      %33913 = stablehlo.select %33909, %33912, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33914 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33915 = stablehlo.dynamic_update_slice %iterArg_537, %33907, %33913, %33914 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33916 = stablehlo.broadcast_in_dim %33474, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33917 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33918 = stablehlo.compare  LT, %iterArg_488, %33917,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33919 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33920 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33921 = stablehlo.add %33919, %33920 : tensor<i32> loc(#loc1845)
      %33922 = stablehlo.select %33918, %33921, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33923 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33924 = stablehlo.dynamic_update_slice %iterArg_538, %33916, %33922, %33923 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33925 = stablehlo.broadcast_in_dim %33489, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc1848)
      %33926 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33927 = stablehlo.compare  LT, %iterArg_488, %33926,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33928 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33929 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33930 = stablehlo.add %33928, %33929 : tensor<i32> loc(#loc1845)
      %33931 = stablehlo.select %33927, %33930, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33932 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33933 = stablehlo.dynamic_update_slice %iterArg_539, %33925, %33931, %33932 : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1847)
      %33934 = stablehlo.broadcast_in_dim %iterArg_489, dims = [1, 2, 3, 4] : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1842)
      %33935 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33936 = stablehlo.compare  LT, %iterArg_488, %33935,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33937 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33938 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33939 = stablehlo.add %33937, %33938 : tensor<i32> loc(#loc1845)
      %33940 = stablehlo.select %33936, %33939, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33941 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33942 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33943 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33944 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33945 = stablehlo.dynamic_update_slice %iterArg_540, %33934, %33940, %33941, %33942, %33943, %33944 : (tensor<1x1x1x2048x2048xbf16>, tensor<1x1x1x2048x2048xbf16>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1847)
      %33946 = stablehlo.broadcast_in_dim %iterArg_490, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1849)
      %33947 = stablehlo.constant dense<0> : tensor<i32> loc(#loc302)
      %33948 = stablehlo.compare  LT, %iterArg_488, %33947,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1843)
      %33949 = stablehlo.convert %iterArg_488 : tensor<i32> loc(#loc1844)
      %33950 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33951 = stablehlo.add %33949, %33950 : tensor<i32> loc(#loc1845)
      %33952 = stablehlo.select %33948, %33951, %iterArg_488 : tensor<i1>, tensor<i32> loc(#loc1846)
      %33953 = stablehlo.dynamic_update_slice %iterArg_541, %33946, %33952 : (tensor<1xi32>, tensor<1xi32>, tensor<i32>) -> tensor<1xi32> loc(#loc1847)
      %33954 = stablehlo.constant dense<1> : tensor<i32> loc(#loc302)
      %33955 = stablehlo.add %iterArg_488, %33954 : tensor<i32> loc(#loc1845)
      stablehlo.return %iterArg, %iterArg_0, %iterArg_1, %iterArg_2, %iterArg_3, %iterArg_4, %iterArg_5, %iterArg_6, %iterArg_7, %iterArg_8, %iterArg_9, %iterArg_10, %iterArg_11, %iterArg_12, %iterArg_13, %iterArg_14, %iterArg_15, %iterArg_16, %iterArg_17, %iterArg_18, %iterArg_19, %iterArg_20, %iterArg_21, %iterArg_22, %iterArg_23, %iterArg_24, %iterArg_25, %iterArg_26, %iterArg_27, %iterArg_28, %iterArg_29, %iterArg_30, %iterArg_31, %iterArg_32, %iterArg_33, %iterArg_34, %iterArg_35, %iterArg_36, %iterArg_37, %iterArg_38, %iterArg_39, %iterArg_40, %iterArg_41, %iterArg_42, %iterArg_43, %iterArg_44, %iterArg_45, %iterArg_46, %iterArg_47, %iterArg_48, %iterArg_49, %iterArg_50, %iterArg_51, %iterArg_52, %iterArg_53, %iterArg_54, %iterArg_55, %iterArg_56, %iterArg_57, %iterArg_58, %iterArg_59, %iterArg_60, %iterArg_61, %iterArg_62, %iterArg_63, %iterArg_64, %iterArg_65, %iterArg_66, %iterArg_67, %iterArg_68, %iterArg_69, %iterArg_70, %iterArg_71, %iterArg_72, %iterArg_73, %iterArg_74, %iterArg_75, %iterArg_76, %iterArg_77, %iterArg_78, %iterArg_79, %iterArg_80, %iterArg_81, %iterArg_82, %iterArg_83, %iterArg_84, %iterArg_85, %iterArg_86, %iterArg_87, %iterArg_88, %iterArg_89, %iterArg_90, %iterArg_91, %iterArg_92, %iterArg_93, %iterArg_94, %iterArg_95, %iterArg_96, %iterArg_97, %iterArg_98, %iterArg_99, %iterArg_100, %iterArg_101, %iterArg_102, %iterArg_103, %iterArg_104, %iterArg_105, %iterArg_106, %iterArg_107, %iterArg_108, %iterArg_109, %iterArg_110, %iterArg_111, %iterArg_112, %iterArg_113, %iterArg_114, %iterArg_115, %iterArg_116, %iterArg_117, %iterArg_118, %iterArg_119, %iterArg_120, %iterArg_121, %iterArg_122, %iterArg_123, %iterArg_124, %iterArg_125, %iterArg_126, %iterArg_127, %iterArg_128, %iterArg_129, %iterArg_130, %iterArg_131, %iterArg_132, %iterArg_133, %iterArg_134, %iterArg_135, %iterArg_136, %iterArg_137, %iterArg_138, %iterArg_139, %iterArg_140, %iterArg_141, %iterArg_142, %iterArg_143, %iterArg_144, %iterArg_145, %iterArg_146, %iterArg_147, %iterArg_148, %iterArg_149, %iterArg_150, %iterArg_151, %iterArg_152, %iterArg_153, %iterArg_154, %iterArg_155, %iterArg_156, %iterArg_157, %iterArg_158, %iterArg_159, %iterArg_160, %iterArg_161, %iterArg_162, %iterArg_163, %iterArg_164, %iterArg_165, %iterArg_166, %iterArg_167, %iterArg_168, %iterArg_169, %iterArg_170, %iterArg_171, %iterArg_172, %iterArg_173, %iterArg_174, %iterArg_175, %iterArg_176, %iterArg_177, %iterArg_178, %iterArg_179, %iterArg_180, %iterArg_181, %iterArg_182, %iterArg_183, %iterArg_184, %iterArg_185, %iterArg_186, %iterArg_187, %iterArg_188, %iterArg_189, %iterArg_190, %iterArg_191, %iterArg_192, %iterArg_193, %iterArg_194, %iterArg_195, %iterArg_196, %iterArg_197, %iterArg_198, %iterArg_199, %iterArg_200, %iterArg_201, %iterArg_202, %iterArg_203, %iterArg_204, %iterArg_205, %iterArg_206, %iterArg_207, %iterArg_208, %iterArg_209, %iterArg_210, %iterArg_211, %iterArg_212, %iterArg_213, %iterArg_214, %iterArg_215, %iterArg_216, %iterArg_217, %iterArg_218, %iterArg_219, %iterArg_220, %iterArg_221, %iterArg_222, %iterArg_223, %iterArg_224, %iterArg_225, %iterArg_226, %iterArg_227, %iterArg_228, %iterArg_229, %iterArg_230, %iterArg_231, %iterArg_232, %iterArg_233, %iterArg_234, %iterArg_235, %iterArg_236, %iterArg_237, %iterArg_238, %iterArg_239, %iterArg_240, %iterArg_241, %iterArg_242, %iterArg_243, %iterArg_244, %iterArg_245, %iterArg_246, %iterArg_247, %iterArg_248, %iterArg_249, %iterArg_250, %iterArg_251, %iterArg_252, %iterArg_253, %iterArg_254, %iterArg_255, %iterArg_256, %iterArg_257, %iterArg_258, %iterArg_259, %iterArg_260, %iterArg_261, %iterArg_262, %iterArg_263, %iterArg_264, %iterArg_265, %iterArg_266, %iterArg_267, %iterArg_268, %iterArg_269, %iterArg_270, %iterArg_271, %iterArg_272, %iterArg_273, %iterArg_274, %iterArg_275, %iterArg_276, %iterArg_277, %iterArg_278, %iterArg_279, %iterArg_280, %iterArg_281, %iterArg_282, %iterArg_283, %iterArg_284, %iterArg_285, %iterArg_286, %iterArg_287, %iterArg_288, %iterArg_289, %iterArg_290, %iterArg_291, %iterArg_292, %iterArg_293, %iterArg_294, %iterArg_295, %iterArg_296, %iterArg_297, %iterArg_298, %iterArg_299, %iterArg_300, %iterArg_301, %iterArg_302, %iterArg_303, %iterArg_304, %iterArg_305, %iterArg_306, %iterArg_307, %iterArg_308, %iterArg_309, %iterArg_310, %iterArg_311, %iterArg_312, %iterArg_313, %iterArg_314, %iterArg_315, %iterArg_316, %iterArg_317, %iterArg_318, %iterArg_319, %iterArg_320, %iterArg_321, %iterArg_322, %iterArg_323, %iterArg_324, %iterArg_325, %iterArg_326, %iterArg_327, %iterArg_328, %iterArg_329, %iterArg_330, %iterArg_331, %iterArg_332, %iterArg_333, %iterArg_334, %iterArg_335, %iterArg_336, %iterArg_337, %iterArg_338, %iterArg_339, %iterArg_340, %iterArg_341, %iterArg_342, %iterArg_343, %iterArg_344, %iterArg_345, %iterArg_346, %iterArg_347, %iterArg_348, %iterArg_349, %iterArg_350, %iterArg_351, %iterArg_352, %iterArg_353, %iterArg_354, %iterArg_355, %iterArg_356, %iterArg_357, %iterArg_358, %iterArg_359, %iterArg_360, %iterArg_361, %iterArg_362, %iterArg_363, %iterArg_364, %iterArg_365, %iterArg_366, %iterArg_367, %iterArg_368, %iterArg_369, %iterArg_370, %iterArg_371, %iterArg_372, %iterArg_373, %iterArg_374, %iterArg_375, %iterArg_376, %iterArg_377, %iterArg_378, %iterArg_379, %iterArg_380, %iterArg_381, %iterArg_382, %iterArg_383, %iterArg_384, %iterArg_385, %iterArg_386, %iterArg_387, %iterArg_388, %iterArg_389, %iterArg_390, %iterArg_391, %iterArg_392, %iterArg_393, %iterArg_394, %iterArg_395, %iterArg_396, %iterArg_397, %iterArg_398, %iterArg_399, %iterArg_400, %iterArg_401, %iterArg_402, %iterArg_403, %iterArg_404, %iterArg_405, %iterArg_406, %iterArg_407, %iterArg_408, %iterArg_409, %iterArg_410, %iterArg_411, %iterArg_412, %iterArg_413, %iterArg_414, %iterArg_415, %iterArg_416, %iterArg_417, %iterArg_418, %iterArg_419, %iterArg_420, %iterArg_421, %iterArg_422, %iterArg_423, %iterArg_424, %iterArg_425, %iterArg_426, %iterArg_427, %iterArg_428, %iterArg_429, %iterArg_430, %iterArg_431, %iterArg_432, %iterArg_433, %iterArg_434, %iterArg_435, %iterArg_436, %iterArg_437, %iterArg_438, %iterArg_439, %iterArg_440, %iterArg_441, %iterArg_442, %iterArg_443, %iterArg_444, %iterArg_445, %iterArg_446, %iterArg_447, %iterArg_448, %iterArg_449, %iterArg_450, %iterArg_451, %iterArg_452, %iterArg_453, %iterArg_454, %iterArg_455, %iterArg_456, %iterArg_457, %iterArg_458, %iterArg_459, %iterArg_460, %iterArg_461, %iterArg_462, %iterArg_463, %iterArg_464, %iterArg_465, %iterArg_466, %iterArg_467, %iterArg_468, %iterArg_469, %iterArg_470, %iterArg_471, %iterArg_472, %iterArg_473, %iterArg_474, %iterArg_475, %iterArg_476, %iterArg_477, %iterArg_478, %iterArg_479, %iterArg_480, %iterArg_481, %iterArg_482, %iterArg_483, %iterArg_484, %iterArg_485, %iterArg_486, %iterArg_487, %33955, %32799, %32801, %33501, %33510, %33519, %33528, %33537, %33546, %33555, %33564, %33573, %33582, %33591, %33600, %33609, %33618, %33627, %33636, %33645, %33654, %33663, %33672, %33681, %33690, %33699, %33708, %33717, %33726, %33735, %33744, %33753, %33762, %33771, %33780, %33789, %33798, %33807, %33816, %33825, %33834, %33843, %33852, %33861, %33870, %33879, %33888, %33897, %33906, %33915, %33924, %33933, %33945, %33953 : tensor<1xi32>, tensor<i32>, tensor<i1>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1x1x2048x2048xbf16>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<i1>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1x2048x2048xi1>, tensor<1xi32>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1x2048xi32>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x64xbf16>, tensor<bf16>, tensor<1x1x32x2048x2048xf32>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x1x2048x2048xf32>, tensor<1xf32>, tensor<1xf32>, tensor<i32>, tensor<1x1x2048x2048xbf16>, tensor<i32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1x1x2048x2048xbf16>, tensor<1xi32> loc(#loc302)
    } loc(#loc302)
    %677 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %678 = stablehlo.reduce(%676#493 init: %677) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %679 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %680 = stablehlo.broadcast_in_dim %679, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %681 = stablehlo.divide %678, %680 : tensor<1xf32> loc(#loc1851)
    %682 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %683 = stablehlo.reduce(%676#494 init: %682) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %684 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %685 = stablehlo.broadcast_in_dim %684, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %686 = stablehlo.divide %683, %685 : tensor<1xf32> loc(#loc1851)
    %687 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %688 = stablehlo.reduce(%676#495 init: %687) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %689 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %690 = stablehlo.broadcast_in_dim %689, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %691 = stablehlo.divide %688, %690 : tensor<1xf32> loc(#loc1851)
    %692 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %693 = stablehlo.reduce(%676#496 init: %692) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %694 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %695 = stablehlo.broadcast_in_dim %694, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %696 = stablehlo.divide %693, %695 : tensor<1xf32> loc(#loc1851)
    %697 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %698 = stablehlo.reduce(%676#497 init: %697) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %699 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %700 = stablehlo.broadcast_in_dim %699, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %701 = stablehlo.divide %698, %700 : tensor<1xf32> loc(#loc1851)
    %702 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %703 = stablehlo.reduce(%676#498 init: %702) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %704 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %705 = stablehlo.broadcast_in_dim %704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %706 = stablehlo.divide %703, %705 : tensor<1xf32> loc(#loc1851)
    %707 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %708 = stablehlo.reduce(%676#499 init: %707) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %709 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %710 = stablehlo.broadcast_in_dim %709, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %711 = stablehlo.divide %708, %710 : tensor<1xf32> loc(#loc1851)
    %712 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %713 = stablehlo.reduce(%676#500 init: %712) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %714 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %715 = stablehlo.broadcast_in_dim %714, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %716 = stablehlo.divide %713, %715 : tensor<1xf32> loc(#loc1851)
    %717 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %718 = stablehlo.reduce(%676#501 init: %717) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %719 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %720 = stablehlo.broadcast_in_dim %719, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %721 = stablehlo.divide %718, %720 : tensor<1xf32> loc(#loc1851)
    %722 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %723 = stablehlo.reduce(%676#502 init: %722) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %724 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %725 = stablehlo.broadcast_in_dim %724, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %726 = stablehlo.divide %723, %725 : tensor<1xf32> loc(#loc1851)
    %727 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %728 = stablehlo.reduce(%676#503 init: %727) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %729 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %730 = stablehlo.broadcast_in_dim %729, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %731 = stablehlo.divide %728, %730 : tensor<1xf32> loc(#loc1851)
    %732 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %733 = stablehlo.reduce(%676#504 init: %732) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %734 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %735 = stablehlo.broadcast_in_dim %734, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %736 = stablehlo.divide %733, %735 : tensor<1xf32> loc(#loc1851)
    %737 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %738 = stablehlo.reduce(%676#505 init: %737) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %739 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %740 = stablehlo.broadcast_in_dim %739, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %741 = stablehlo.divide %738, %740 : tensor<1xf32> loc(#loc1851)
    %742 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %743 = stablehlo.reduce(%676#506 init: %742) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %744 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %745 = stablehlo.broadcast_in_dim %744, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %746 = stablehlo.divide %743, %745 : tensor<1xf32> loc(#loc1851)
    %747 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %748 = stablehlo.reduce(%676#507 init: %747) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %749 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %750 = stablehlo.broadcast_in_dim %749, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %751 = stablehlo.divide %748, %750 : tensor<1xf32> loc(#loc1851)
    %752 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %753 = stablehlo.reduce(%676#508 init: %752) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %754 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %755 = stablehlo.broadcast_in_dim %754, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %756 = stablehlo.divide %753, %755 : tensor<1xf32> loc(#loc1851)
    %757 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %758 = stablehlo.reduce(%676#509 init: %757) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %759 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %760 = stablehlo.broadcast_in_dim %759, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %761 = stablehlo.divide %758, %760 : tensor<1xf32> loc(#loc1851)
    %762 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %763 = stablehlo.reduce(%676#510 init: %762) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %764 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %765 = stablehlo.broadcast_in_dim %764, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %766 = stablehlo.divide %763, %765 : tensor<1xf32> loc(#loc1851)
    %767 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %768 = stablehlo.reduce(%676#511 init: %767) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %769 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %770 = stablehlo.broadcast_in_dim %769, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %771 = stablehlo.divide %768, %770 : tensor<1xf32> loc(#loc1851)
    %772 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %773 = stablehlo.reduce(%676#512 init: %772) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %774 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %775 = stablehlo.broadcast_in_dim %774, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %776 = stablehlo.divide %773, %775 : tensor<1xf32> loc(#loc1851)
    %777 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %778 = stablehlo.reduce(%676#513 init: %777) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %779 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %780 = stablehlo.broadcast_in_dim %779, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %781 = stablehlo.divide %778, %780 : tensor<1xf32> loc(#loc1851)
    %782 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %783 = stablehlo.reduce(%676#514 init: %782) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %784 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %785 = stablehlo.broadcast_in_dim %784, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %786 = stablehlo.divide %783, %785 : tensor<1xf32> loc(#loc1851)
    %787 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %788 = stablehlo.reduce(%676#515 init: %787) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %789 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %790 = stablehlo.broadcast_in_dim %789, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %791 = stablehlo.divide %788, %790 : tensor<1xf32> loc(#loc1851)
    %792 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %793 = stablehlo.reduce(%676#516 init: %792) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %794 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %795 = stablehlo.broadcast_in_dim %794, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %796 = stablehlo.divide %793, %795 : tensor<1xf32> loc(#loc1851)
    %797 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %798 = stablehlo.reduce(%676#517 init: %797) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %799 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %800 = stablehlo.broadcast_in_dim %799, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %801 = stablehlo.divide %798, %800 : tensor<1xf32> loc(#loc1851)
    %802 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %803 = stablehlo.reduce(%676#518 init: %802) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %804 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %805 = stablehlo.broadcast_in_dim %804, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %806 = stablehlo.divide %803, %805 : tensor<1xf32> loc(#loc1851)
    %807 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %808 = stablehlo.reduce(%676#519 init: %807) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %809 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %810 = stablehlo.broadcast_in_dim %809, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %811 = stablehlo.divide %808, %810 : tensor<1xf32> loc(#loc1851)
    %812 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %813 = stablehlo.reduce(%676#520 init: %812) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %814 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %815 = stablehlo.broadcast_in_dim %814, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %816 = stablehlo.divide %813, %815 : tensor<1xf32> loc(#loc1851)
    %817 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %818 = stablehlo.reduce(%676#521 init: %817) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %819 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %820 = stablehlo.broadcast_in_dim %819, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %821 = stablehlo.divide %818, %820 : tensor<1xf32> loc(#loc1851)
    %822 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %823 = stablehlo.reduce(%676#522 init: %822) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %824 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %825 = stablehlo.broadcast_in_dim %824, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %826 = stablehlo.divide %823, %825 : tensor<1xf32> loc(#loc1851)
    %827 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %828 = stablehlo.reduce(%676#523 init: %827) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %829 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %830 = stablehlo.broadcast_in_dim %829, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %831 = stablehlo.divide %828, %830 : tensor<1xf32> loc(#loc1851)
    %832 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %833 = stablehlo.reduce(%676#524 init: %832) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %834 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %835 = stablehlo.broadcast_in_dim %834, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %836 = stablehlo.divide %833, %835 : tensor<1xf32> loc(#loc1851)
    %837 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %838 = stablehlo.reduce(%676#525 init: %837) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %839 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %840 = stablehlo.broadcast_in_dim %839, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %841 = stablehlo.divide %838, %840 : tensor<1xf32> loc(#loc1851)
    %842 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %843 = stablehlo.reduce(%676#526 init: %842) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %844 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %845 = stablehlo.broadcast_in_dim %844, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %846 = stablehlo.divide %843, %845 : tensor<1xf32> loc(#loc1851)
    %847 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %848 = stablehlo.reduce(%676#527 init: %847) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %849 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %850 = stablehlo.broadcast_in_dim %849, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %851 = stablehlo.divide %848, %850 : tensor<1xf32> loc(#loc1851)
    %852 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %853 = stablehlo.reduce(%676#528 init: %852) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %854 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %855 = stablehlo.broadcast_in_dim %854, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %856 = stablehlo.divide %853, %855 : tensor<1xf32> loc(#loc1851)
    %857 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %858 = stablehlo.reduce(%676#529 init: %857) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %859 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %860 = stablehlo.broadcast_in_dim %859, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %861 = stablehlo.divide %858, %860 : tensor<1xf32> loc(#loc1851)
    %862 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %863 = stablehlo.reduce(%676#530 init: %862) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %864 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %865 = stablehlo.broadcast_in_dim %864, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %866 = stablehlo.divide %863, %865 : tensor<1xf32> loc(#loc1851)
    %867 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %868 = stablehlo.reduce(%676#531 init: %867) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %869 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %870 = stablehlo.broadcast_in_dim %869, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %871 = stablehlo.divide %868, %870 : tensor<1xf32> loc(#loc1851)
    %872 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %873 = stablehlo.reduce(%676#532 init: %872) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %874 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %875 = stablehlo.broadcast_in_dim %874, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %876 = stablehlo.divide %873, %875 : tensor<1xf32> loc(#loc1851)
    %877 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %878 = stablehlo.reduce(%676#533 init: %877) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %879 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %880 = stablehlo.broadcast_in_dim %879, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %881 = stablehlo.divide %878, %880 : tensor<1xf32> loc(#loc1851)
    %882 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %883 = stablehlo.reduce(%676#534 init: %882) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %884 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %885 = stablehlo.broadcast_in_dim %884, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %886 = stablehlo.divide %883, %885 : tensor<1xf32> loc(#loc1851)
    %887 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %888 = stablehlo.reduce(%676#535 init: %887) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %889 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %890 = stablehlo.broadcast_in_dim %889, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %891 = stablehlo.divide %888, %890 : tensor<1xf32> loc(#loc1851)
    %892 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %893 = stablehlo.reduce(%676#536 init: %892) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %894 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %895 = stablehlo.broadcast_in_dim %894, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %896 = stablehlo.divide %893, %895 : tensor<1xf32> loc(#loc1851)
    %897 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %898 = stablehlo.reduce(%676#537 init: %897) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %899 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %900 = stablehlo.broadcast_in_dim %899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %901 = stablehlo.divide %898, %900 : tensor<1xf32> loc(#loc1851)
    %902 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %903 = stablehlo.reduce(%676#538 init: %902) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %904 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %905 = stablehlo.broadcast_in_dim %904, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %906 = stablehlo.divide %903, %905 : tensor<1xf32> loc(#loc1851)
    %907 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %908 = stablehlo.reduce(%676#539 init: %907) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %909 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %910 = stablehlo.broadcast_in_dim %909, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %911 = stablehlo.divide %908, %910 : tensor<1xf32> loc(#loc1851)
    %912 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1850)
    %913 = stablehlo.reduce(%676#540 init: %912) across dimensions = [0] : (tensor<1x1xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1850)
      stablehlo.return %29929 : tensor<f32> loc(#loc1850)
    } loc(#loc1850)
    %914 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %915 = stablehlo.broadcast_in_dim %914, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc1851)
    %916 = stablehlo.divide %913, %915 : tensor<1xf32> loc(#loc1851)
    %917 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %918 = stablehlo.dynamic_slice %681, %917, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %919 = stablehlo.reshape %918 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %920 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %921 = stablehlo.dynamic_slice %686, %920, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %922 = stablehlo.reshape %921 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %923 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %924 = stablehlo.dynamic_slice %691, %923, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %925 = stablehlo.reshape %924 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %926 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %927 = stablehlo.dynamic_slice %696, %926, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %928 = stablehlo.reshape %927 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %929 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %930 = stablehlo.dynamic_slice %701, %929, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %931 = stablehlo.reshape %930 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %932 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %933 = stablehlo.dynamic_slice %706, %932, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %934 = stablehlo.reshape %933 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %935 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %936 = stablehlo.dynamic_slice %711, %935, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %937 = stablehlo.reshape %936 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %938 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %939 = stablehlo.dynamic_slice %716, %938, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %940 = stablehlo.reshape %939 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %941 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %942 = stablehlo.dynamic_slice %721, %941, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %943 = stablehlo.reshape %942 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %944 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %945 = stablehlo.dynamic_slice %726, %944, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %946 = stablehlo.reshape %945 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %947 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %948 = stablehlo.dynamic_slice %731, %947, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %949 = stablehlo.reshape %948 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %950 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %951 = stablehlo.dynamic_slice %736, %950, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %952 = stablehlo.reshape %951 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %953 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %954 = stablehlo.dynamic_slice %741, %953, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %955 = stablehlo.reshape %954 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %956 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %957 = stablehlo.dynamic_slice %746, %956, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %958 = stablehlo.reshape %957 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %959 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %960 = stablehlo.dynamic_slice %751, %959, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %961 = stablehlo.reshape %960 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %962 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %963 = stablehlo.dynamic_slice %756, %962, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %964 = stablehlo.reshape %963 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %965 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %966 = stablehlo.dynamic_slice %761, %965, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %967 = stablehlo.reshape %966 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %968 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %969 = stablehlo.dynamic_slice %766, %968, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %970 = stablehlo.reshape %969 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %971 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %972 = stablehlo.dynamic_slice %771, %971, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %973 = stablehlo.reshape %972 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %974 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %975 = stablehlo.dynamic_slice %776, %974, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %976 = stablehlo.reshape %975 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %977 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %978 = stablehlo.dynamic_slice %781, %977, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %979 = stablehlo.reshape %978 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %980 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %981 = stablehlo.dynamic_slice %786, %980, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %982 = stablehlo.reshape %981 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %983 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %984 = stablehlo.dynamic_slice %791, %983, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %985 = stablehlo.reshape %984 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %986 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %987 = stablehlo.dynamic_slice %796, %986, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %988 = stablehlo.reshape %987 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %989 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %990 = stablehlo.dynamic_slice %801, %989, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %991 = stablehlo.reshape %990 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %992 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %993 = stablehlo.dynamic_slice %806, %992, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %994 = stablehlo.reshape %993 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %995 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %996 = stablehlo.dynamic_slice %811, %995, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %997 = stablehlo.reshape %996 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %998 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %999 = stablehlo.dynamic_slice %816, %998, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1000 = stablehlo.reshape %999 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1001 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1002 = stablehlo.dynamic_slice %821, %1001, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1003 = stablehlo.reshape %1002 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1004 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1005 = stablehlo.dynamic_slice %826, %1004, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1006 = stablehlo.reshape %1005 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1007 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1008 = stablehlo.dynamic_slice %831, %1007, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1009 = stablehlo.reshape %1008 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1010 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1011 = stablehlo.dynamic_slice %836, %1010, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1012 = stablehlo.reshape %1011 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1013 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1014 = stablehlo.dynamic_slice %841, %1013, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1015 = stablehlo.reshape %1014 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1016 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1017 = stablehlo.dynamic_slice %846, %1016, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1018 = stablehlo.reshape %1017 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1019 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1020 = stablehlo.dynamic_slice %851, %1019, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1021 = stablehlo.reshape %1020 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1022 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1023 = stablehlo.dynamic_slice %856, %1022, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1024 = stablehlo.reshape %1023 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1025 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1026 = stablehlo.dynamic_slice %861, %1025, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1027 = stablehlo.reshape %1026 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1028 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1029 = stablehlo.dynamic_slice %866, %1028, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1030 = stablehlo.reshape %1029 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1031 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1032 = stablehlo.dynamic_slice %871, %1031, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1033 = stablehlo.reshape %1032 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1034 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1035 = stablehlo.dynamic_slice %876, %1034, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1036 = stablehlo.reshape %1035 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1037 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1038 = stablehlo.dynamic_slice %881, %1037, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1039 = stablehlo.reshape %1038 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1040 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1041 = stablehlo.dynamic_slice %886, %1040, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1042 = stablehlo.reshape %1041 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1043 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1044 = stablehlo.dynamic_slice %891, %1043, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1045 = stablehlo.reshape %1044 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1046 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1047 = stablehlo.dynamic_slice %896, %1046, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1048 = stablehlo.reshape %1047 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1049 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1050 = stablehlo.dynamic_slice %901, %1049, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1051 = stablehlo.reshape %1050 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1052 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1053 = stablehlo.dynamic_slice %906, %1052, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1054 = stablehlo.reshape %1053 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1055 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1056 = stablehlo.dynamic_slice %911, %1055, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1057 = stablehlo.reshape %1056 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1058 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1059 = stablehlo.dynamic_slice %916, %1058, sizes = [1] : (tensor<1xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc1852)
    %1060 = stablehlo.reshape %1059 : (tensor<1xf32>) -> tensor<f32> loc(#loc1853)
    %1061 = "stablehlo.slice"(%676#492) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1854)
    %1062 = stablehlo.reshape %1061 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1855)
    %1063 = stablehlo.reshape %1062 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1856)
    %1064 = stablehlo.transpose %1063, dims = [0, 2, 1, 3, 4] : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1857)
    %1065 = stablehlo.reshape %1064 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1858)
    %1066 = arith.select %ptrue, %1065, %1065 : tensor<1x2048x2048xbf16>
    %1067 = arith.select %ptrue, %1066, %1066 : tensor<1x2048x2048xbf16>
    %1068 = stablehlo.convert %1067 : (tensor<1x2048x2048xbf16>) -> tensor<1x2048x2048xf32> loc(#loc1859)
    %1069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1860)
    %1070 = stablehlo.reduce(%1068 init: %1069) across dimensions = [2] : (tensor<1x2048x2048xf32>, tensor<f32>) -> tensor<1x2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1860)
      stablehlo.return %29929 : tensor<f32> loc(#loc1860)
    } loc(#loc1860)
    %1071 = stablehlo.broadcast_in_dim %1070, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x1xf32> loc(#loc1861)
    %1072 = stablehlo.convert %1071 : (tensor<1x2048x1xf32>) -> tensor<1x2048x1xbf16> loc(#loc1862)
    %1073 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc)
    %1074 = stablehlo.broadcast_in_dim %1073, dims = [] : (tensor<bf16>) -> tensor<1x2048x1xbf16> loc(#loc1863)
    %1075 = stablehlo.divide %1072, %1074 : tensor<1x2048x1xbf16> loc(#loc1863)
    %1076 = stablehlo.broadcast_in_dim %1075, dims = [0, 1, 2] : (tensor<1x2048x1xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1864)
    %1077 = stablehlo.subtract %1067, %1076 : tensor<1x2048x2048xbf16> loc(#loc1864)
    %1078 = stablehlo.multiply %1077, %1077 : tensor<1x2048x2048xbf16> loc(#loc1865)
    %1079 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc)
    %1080 = stablehlo.broadcast_in_dim %1079, dims = [] : (tensor<bf16>) -> tensor<1x2048x2048xbf16> loc(#loc1865)
    %1081 = stablehlo.multiply %1080, %1077 : tensor<1x2048x2048xbf16> loc(#loc1865)
    %1082 = stablehlo.convert %1078 : (tensor<1x2048x2048xbf16>) -> tensor<1x2048x2048xf32> loc(#loc1859)
    %1083 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1860)
    %1084 = stablehlo.reduce(%1082 init: %1083) across dimensions = [2] : (tensor<1x2048x2048xf32>, tensor<f32>) -> tensor<1x2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1860)
      stablehlo.return %29929 : tensor<f32> loc(#loc1860)
    } loc(#loc1860)
    %1085 = stablehlo.broadcast_in_dim %1084, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x1xf32> loc(#loc1861)
    %1086 = stablehlo.convert %1085 : (tensor<1x2048x1xf32>) -> tensor<1x2048x1xbf16> loc(#loc1862)
    %1087 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc)
    %1088 = stablehlo.broadcast_in_dim %1087, dims = [] : (tensor<bf16>) -> tensor<1x2048x1xbf16> loc(#loc1863)
    %1089 = stablehlo.divide %1086, %1088 : tensor<1x2048x1xbf16> loc(#loc1863)
    %1090 = stablehlo.broadcast_in_dim %1075, dims = [0, 1, 2] : (tensor<1x2048x1xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1864)
    %1091 = stablehlo.subtract %1067, %1090 : tensor<1x2048x2048xbf16> loc(#loc1864)
    %1092 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc)
    %1093 = stablehlo.broadcast_in_dim %1092, dims = [] : (tensor<bf16>) -> tensor<1x2048x1xbf16> loc(#loc1866)
    %1094 = stablehlo.add %1089, %1093 : tensor<1x2048x1xbf16> loc(#loc1866)
    %1095 = stablehlo.rsqrt %1094 : tensor<1x2048x1xbf16> loc(#loc1867)
    %1096 = stablehlo.divide %1095, %1094 : tensor<1x2048x1xbf16> loc(#loc1868)
    %1097 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc)
    %1098 = stablehlo.broadcast_in_dim %1097, dims = [] : (tensor<bf16>) -> tensor<1x2048x1xbf16> loc(#loc1869)
    %1099 = stablehlo.multiply %1098, %1096 : tensor<1x2048x1xbf16> loc(#loc1869)
    %1100 = stablehlo.broadcast_in_dim %1095, dims = [0, 1, 2] : (tensor<1x2048x1xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1869)
    %1101 = stablehlo.multiply %1091, %1100 : tensor<1x2048x2048xbf16> loc(#loc1869)
    %1102 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc)
    %1103 = stablehlo.broadcast_in_dim %1102, dims = [] : (tensor<bf16>) -> tensor<2048xbf16> loc(#loc1870)
    %1104 = stablehlo.add %1103, %13 : tensor<2048xbf16> loc(#loc1870)
    %1105 = stablehlo.broadcast_in_dim %1104, dims = [2] : (tensor<2048xbf16>) -> tensor<1x1x2048xbf16> loc(#loc1871)
    %1106 = stablehlo.broadcast_in_dim %1105, dims = [0, 1, 2] : (tensor<1x1x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1872)
    %1107 = stablehlo.multiply %1101, %1106 : tensor<1x2048x2048xbf16> loc(#loc1872)
    %1108 = stablehlo.broadcast_in_dim %12, dims = [2] : (tensor<2048xbf16>) -> tensor<1x1x2048xbf16> loc(#loc1873)
    %1109 = stablehlo.broadcast_in_dim %1108, dims = [0, 1, 2] : (tensor<1x1x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1874)
    %1110 = stablehlo.add %1107, %1109 : tensor<1x2048x2048xbf16> loc(#loc1874)
    %1111 = stablehlo.broadcast_in_dim %7, dims = [0, 1] : (tensor<1x2048xi32>) -> tensor<1x2048x1xi32> loc(#loc1875)
    %1112 = stablehlo.broadcast_in_dim %11, dims = [0, 1] : (tensor<1x2048xi32>) -> tensor<1x2048x1xi32> loc(#loc1876)
    %1113 = call @_einsum_455(%1110, %16) : (tensor<1x2048x2048xbf16>, tensor<2048x51200xbf16>) -> tensor<1x2048x51200xbf16> loc(#loc1877)
    %1114 = arith.select %ptrue, %1113, %1113 : tensor<1x2048x51200xbf16>
    %1115 = stablehlo.broadcast_in_dim %15, dims = [2] : (tensor<51200xbf16>) -> tensor<1x1x51200xbf16> loc(#loc1879)
    %1116 = stablehlo.broadcast_in_dim %1115, dims = [0, 1, 2] : (tensor<1x1x51200xbf16>) -> tensor<1x2048x51200xbf16> loc(#loc1880)
    %1117 = stablehlo.add %1114, %1116 : tensor<1x2048x51200xbf16> loc(#loc1880)
    %1118 = stablehlo.constant dense<3.000000e+01> : tensor<bf16> loc(#loc)
    %1119 = stablehlo.broadcast_in_dim %1118, dims = [] : (tensor<bf16>) -> tensor<1x2048x51200xbf16> loc(#loc1881)
    %1120 = stablehlo.divide %1117, %1119 : tensor<1x2048x51200xbf16> loc(#loc1881)
    %1121 = stablehlo.tanh %1120 : tensor<1x2048x51200xbf16> loc(#loc1882)
    %1122 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc)
    %1123 = stablehlo.broadcast_in_dim %1122, dims = [] : (tensor<bf16>) -> tensor<1x2048x51200xbf16> loc(#loc1883)
    %1124 = stablehlo.subtract %1123, %1121 : tensor<1x2048x51200xbf16> loc(#loc1883)
    %1125 = stablehlo.constant dense<3.000000e+01> : tensor<bf16> loc(#loc)
    %1126 = stablehlo.broadcast_in_dim %1125, dims = [] : (tensor<bf16>) -> tensor<1x2048x51200xbf16> loc(#loc1884)
    %1127 = stablehlo.multiply %1126, %1121 : tensor<1x2048x51200xbf16> loc(#loc1884)
    %1128 = stablehlo.convert %1127 : (tensor<1x2048x51200xbf16>) -> tensor<1x2048x51200xf32> loc(#loc1885)
    %1129:3 = call @log_softmax(%1128) : (tensor<1x2048x51200xf32>) -> (tensor<1x2048x51200xf32>, tensor<1x2048x51200xf32>, tensor<1x2048x1xf32>) loc(#loc1886)
    %1130 = stablehlo.reshape %1111 : (tensor<1x2048x1xi32>) -> tensor<1x2048xi32> loc(#loc1887)
    %1131 = call @_one_hot_456(%1130) : (tensor<1x2048xi32>) -> tensor<1x2048x51200xf32> loc(#loc1888)
    %1132 = stablehlo.multiply %1129#0, %1131 : tensor<1x2048x51200xf32> loc(#loc1889)
    %1133 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1890)
    %1134 = stablehlo.reduce(%1132 init: %1133) across dimensions = [2] : (tensor<1x2048x51200xf32>, tensor<f32>) -> tensor<1x2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1890)
      stablehlo.return %29929 : tensor<f32> loc(#loc1890)
    } loc(#loc1890)
    %1135 = stablehlo.negate %1134 : tensor<1x2048xf32> loc(#loc1891)
    %1136 = call @argmax(%1128) : (tensor<1x2048x51200xf32>) -> tensor<1x2048xi32> loc(#loc1892)
    %1137 = stablehlo.broadcast_in_dim %1135, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x1xf32> loc(#loc1893)
    %1138 = stablehlo.convert %1112 : (tensor<1x2048x1xi32>) -> tensor<1x2048x1xf32> loc(#loc1894)
    %1139 = stablehlo.multiply %1137, %1138 : tensor<1x2048x1xf32> loc(#loc1895)
    %1140 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1896)
    %1141 = stablehlo.reduce(%1139 init: %1140) across dimensions = [0, 1, 2] : (tensor<1x2048x1xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1896)
      stablehlo.return %29929 : tensor<f32> loc(#loc1896)
    } loc(#loc1896)
    %1142 = stablehlo.convert %1112 : (tensor<1x2048x1xi32>) -> tensor<1x2048x1xf32> loc(#loc1897)
    %1143 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1898)
    %1144 = stablehlo.reduce(%1142 init: %1143) across dimensions = [0, 1, 2] : (tensor<1x2048x1xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1898)
      stablehlo.return %29929 : tensor<f32> loc(#loc1898)
    } loc(#loc1898)
    %1145 = stablehlo.constant dense<9.99999997E-7> : tensor<f32> loc(#loc)
    %1146 = stablehlo.add %1144, %1145 : tensor<f32> loc(#loc1899)
    %1147 = stablehlo.divide %1141, %1146 : tensor<f32> loc(#loc1900)
    %1148 = stablehlo.convert %11 : (tensor<1x2048xi32>) -> tensor<1x2048xf32> loc(#loc1901)
    %1149 = stablehlo.multiply %1135, %1148 : tensor<1x2048xf32> loc(#loc1902)
    %1150 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1903)
    %1151 = stablehlo.reduce(%1149 init: %1150) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1903)
      stablehlo.return %29929 : tensor<f32> loc(#loc1903)
    } loc(#loc1903)
    %1152 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1153 = stablehlo.add %1147, %1152 : tensor<f32> loc(#loc1904)
    %1154 = stablehlo.compare  EQ, %7, %1136,  SIGNED : (tensor<1x2048xi32>, tensor<1x2048xi32>) -> tensor<1x2048xi1> loc(#loc1905)
    %1155 = stablehlo.convert %1154 : (tensor<1x2048xi1>) -> tensor<1x2048xi32> loc(#loc1906)
    %1156 = stablehlo.multiply %1155, %11 : tensor<1x2048xi32> loc(#loc1907)
    %1157 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1908)
    %1158 = stablehlo.reduce(%1156 init: %1157) across dimensions = [0, 1] : (tensor<1x2048xi32>, tensor<i32>) -> tensor<i32>
     reducer(%arg558: tensor<i32> loc(unknown), %arg559: tensor<i32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<i32> loc(#loc1908)
      stablehlo.return %29929 : tensor<i32> loc(#loc1908)
    } loc(#loc1908)
    %1159 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1160 = stablehlo.maximum %1144, %1159 : tensor<f32> loc(#loc1909)
    %1161 = stablehlo.convert %1158 : (tensor<i32>) -> tensor<f32> loc(#loc1910)
    %1162 = stablehlo.divide %1161, %1160 : tensor<f32> loc(#loc1911)
    %1163 = stablehlo.negate %1151 : tensor<1xf32> loc(#loc1912)
    %1164 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1165 = stablehlo.divide %1164, %1146 : tensor<f32> loc(#loc1913)
    %1166 = stablehlo.broadcast_in_dim %1165, dims = [] : (tensor<f32>) -> tensor<1x2048x1xf32> loc(#loc1914)
    %1167 = stablehlo.multiply %1166, %1138 : tensor<1x2048x1xf32> loc(#loc1915)
    %1168 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1916)
    %1169 = stablehlo.reduce(%1167 init: %1168) across dimensions = [0, 2] : (tensor<1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1916)
      stablehlo.return %29929 : tensor<f32> loc(#loc1916)
    } loc(#loc1916)
    %1170 = stablehlo.broadcast_in_dim %1169, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc1917)
    %1171 = stablehlo.negate %1170 : tensor<1x2048xf32> loc(#loc1918)
    %1172 = stablehlo.broadcast_in_dim %1171, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x51200xf32> loc(#loc1919)
    %1173 = stablehlo.multiply %1172, %1131 : tensor<1x2048x51200xf32> loc(#loc1920)
    %1174 = call @log_softmax_457(%1129#1, %1129#2, %1173) : (tensor<1x2048x51200xf32>, tensor<1x2048x1xf32>, tensor<1x2048x51200xf32>) -> tensor<1x2048x51200xf32> loc(#loc1921)
    %1175 = stablehlo.convert %1174 : (tensor<1x2048x51200xf32>) -> tensor<1x2048x51200xbf16> loc(#loc1922)
    %1176 = stablehlo.constant dense<3.000000e+01> : tensor<bf16> loc(#loc)
    %1177 = stablehlo.broadcast_in_dim %1176, dims = [] : (tensor<bf16>) -> tensor<1x2048x51200xbf16> loc(#loc1923)
    %1178 = stablehlo.multiply %1177, %1175 : tensor<1x2048x51200xbf16> loc(#loc1923)
    %1179 = stablehlo.multiply %1178, %1124 : tensor<1x2048x51200xbf16> loc(#loc1923)
    %1180 = stablehlo.multiply %1179, %1121 : tensor<1x2048x51200xbf16> loc(#loc1923)
    %1181 = stablehlo.add %1179, %1180 : tensor<1x2048x51200xbf16> loc(#loc1924)
    %1182 = stablehlo.constant dense<3.000000e+01> : tensor<bf16> loc(#loc)
    %1183 = stablehlo.broadcast_in_dim %1182, dims = [] : (tensor<bf16>) -> tensor<1x2048x51200xbf16> loc(#loc1925)
    %1184 = stablehlo.divide %1181, %1183 : tensor<1x2048x51200xbf16> loc(#loc1925)
    %1185 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1926)
    %1186 = stablehlo.reduce(%1184 init: %1185) across dimensions = [1] : (tensor<1x2048x51200xbf16>, tensor<bf16>) -> tensor<1x51200xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1926)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1926)
    } loc(#loc1926)
    %1187 = stablehlo.reshape %1186 : (tensor<1x51200xbf16>) -> tensor<1x1x51200xbf16> loc(#loc1927)
    %1188 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1928)
    %1189 = stablehlo.reduce(%1187 init: %1188) across dimensions = [0, 1] : (tensor<1x1x51200xbf16>, tensor<bf16>) -> tensor<51200xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1928)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1928)
    } loc(#loc1928)
    %1190 = stablehlo.convert %1189 : (tensor<51200xbf16>) -> tensor<51200xf32> loc(#loc1929)
    %1191 = arith.select %ptrue, %1184, %1184 : tensor<1x2048x51200xbf16>
    %1192:2 = call @_einsum_458(%16, %1110, %1191) : (tensor<2048x51200xbf16>, tensor<1x2048x2048xbf16>, tensor<1x2048x51200xbf16>) -> (tensor<1x2048x2048xbf16>, tensor<2048x51200xbf16>) loc(#loc1931)
    %1193 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1932)
    %1194 = stablehlo.reduce(%1192#0 init: %1193) across dimensions = [1] : (tensor<1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1932)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1932)
    } loc(#loc1932)
    %1195 = stablehlo.reshape %1194 : (tensor<1x2048xbf16>) -> tensor<1x1x2048xbf16> loc(#loc1933)
    %1196 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1934)
    %1197 = stablehlo.reduce(%1195 init: %1196) across dimensions = [0, 1] : (tensor<1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1934)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1934)
    } loc(#loc1934)
    %1198 = stablehlo.convert %1197 : (tensor<2048xbf16>) -> tensor<2048xf32> loc(#loc1929)
    %1199 = stablehlo.multiply %1101, %1192#0 : tensor<1x2048x2048xbf16> loc(#loc1935)
    %1200 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1936)
    %1201 = stablehlo.reduce(%1199 init: %1200) across dimensions = [1] : (tensor<1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1936)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1936)
    } loc(#loc1936)
    %1202 = stablehlo.reshape %1201 : (tensor<1x2048xbf16>) -> tensor<1x1x2048xbf16> loc(#loc1937)
    %1203 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1938)
    %1204 = stablehlo.reduce(%1202 init: %1203) across dimensions = [0, 1] : (tensor<1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1938)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1938)
    } loc(#loc1938)
    %1205 = stablehlo.broadcast_in_dim %1105, dims = [0, 1, 2] : (tensor<1x1x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1935)
    %1206 = stablehlo.multiply %1192#0, %1205 : tensor<1x2048x2048xbf16> loc(#loc1935)
    %1207 = stablehlo.convert %1204 : (tensor<2048xbf16>) -> tensor<2048xf32> loc(#loc1929)
    %1208 = stablehlo.multiply %1091, %1206 : tensor<1x2048x2048xbf16> loc(#loc1939)
    %1209 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1940)
    %1210 = stablehlo.reduce(%1208 init: %1209) across dimensions = [2] : (tensor<1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1940)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1940)
    } loc(#loc1940)
    %1211 = stablehlo.reshape %1210 : (tensor<1x2048xbf16>) -> tensor<1x2048x1xbf16> loc(#loc1941)
    %1212 = stablehlo.broadcast_in_dim %1095, dims = [0, 1, 2] : (tensor<1x2048x1xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc1939)
    %1213 = stablehlo.multiply %1206, %1212 : tensor<1x2048x2048xbf16> loc(#loc1939)
    %1214 = stablehlo.multiply %1211, %1099 : tensor<1x2048x1xbf16> loc(#loc1939)
    %1215 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc)
    %1216 = stablehlo.broadcast_in_dim %1215, dims = [] : (tensor<bf16>) -> tensor<1x2048x1xbf16> loc(#loc1942)
    %1217 = stablehlo.divide %1214, %1216 : tensor<1x2048x1xbf16> loc(#loc1942)
    %1218 = stablehlo.convert %1217 : (tensor<1x2048x1xbf16>) -> tensor<1x2048x1xf32> loc(#loc1943)
    %1219 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1944)
    %1220 = stablehlo.reduce(%1218 init: %1219) across dimensions = [0, 2] : (tensor<1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1944)
      stablehlo.return %29929 : tensor<f32> loc(#loc1944)
    } loc(#loc1944)
    %1221 = stablehlo.broadcast_in_dim %1220, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc1945)
    %1222 = stablehlo.broadcast_in_dim %1221, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x2048xf32> loc(#loc1946)
    %1223 = stablehlo.convert %1222 : (tensor<1x2048x2048xf32>) -> tensor<1x2048x2048xbf16> loc(#loc1947)
    %1224 = stablehlo.multiply %1223, %1081 : tensor<1x2048x2048xbf16> loc(#loc1948)
    %1225 = stablehlo.negate %1224 : tensor<1x2048x2048xbf16> loc(#loc1949)
    %1226 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1950)
    %1227 = stablehlo.reduce(%1225 init: %1226) across dimensions = [2] : (tensor<1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1950)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1950)
    } loc(#loc1950)
    %1228 = stablehlo.reshape %1227 : (tensor<1x2048xbf16>) -> tensor<1x2048x1xbf16> loc(#loc1951)
    %1229 = stablehlo.negate %1213 : tensor<1x2048x2048xbf16> loc(#loc1949)
    %1230 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1950)
    %1231 = stablehlo.reduce(%1229 init: %1230) across dimensions = [2] : (tensor<1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc1950)
      stablehlo.return %29929 : tensor<bf16> loc(#loc1950)
    } loc(#loc1950)
    %1232 = stablehlo.reshape %1231 : (tensor<1x2048xbf16>) -> tensor<1x2048x1xbf16> loc(#loc1951)
    %1233 = stablehlo.add %1224, %1213 : tensor<1x2048x2048xbf16> loc(#loc1952)
    %1234 = stablehlo.add %1228, %1232 : tensor<1x2048x1xbf16> loc(#loc1952)
    %1235 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc)
    %1236 = stablehlo.broadcast_in_dim %1235, dims = [] : (tensor<bf16>) -> tensor<1x2048x1xbf16> loc(#loc1942)
    %1237 = stablehlo.divide %1234, %1236 : tensor<1x2048x1xbf16> loc(#loc1942)
    %1238 = stablehlo.convert %1237 : (tensor<1x2048x1xbf16>) -> tensor<1x2048x1xf32> loc(#loc1943)
    %1239 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1944)
    %1240 = stablehlo.reduce(%1238 init: %1239) across dimensions = [0, 2] : (tensor<1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc1944)
      stablehlo.return %29929 : tensor<f32> loc(#loc1944)
    } loc(#loc1944)
    %1241 = stablehlo.broadcast_in_dim %1240, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc1945)
    %1242 = stablehlo.broadcast_in_dim %1241, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x2048xf32> loc(#loc1946)
    %1243 = stablehlo.convert %1242 : (tensor<1x2048x2048xf32>) -> tensor<1x2048x2048xbf16> loc(#loc1947)
    %1244 = stablehlo.add %1233, %1243 : tensor<1x2048x2048xbf16> loc(#loc1953)
    %1245 = arith.select %ptrue, %1244, %1244 : tensor<1x2048x2048xbf16>
    %1246 = arith.select %ptrue, %1245, %1245 : tensor<1x2048x2048xbf16>
    %1247 = stablehlo.reshape %1246 : (tensor<1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1955)
    %1248 = stablehlo.transpose %1247, dims = [0, 2, 1, 3, 4] : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1956)
    %1249 = stablehlo.reshape %1248 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1957)
    %1250 = stablehlo.broadcast_in_dim %1249, dims = [0, 2, 3, 4] : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1958)
    %1251 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1252 = stablehlo.pad %1250, %1251, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1959)
    %1253 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1254 = stablehlo.broadcast_in_dim %1253, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1960)
    %1255 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1256 = stablehlo.broadcast_in_dim %1255, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1257 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1258 = stablehlo.broadcast_in_dim %1257, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1259 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1260 = stablehlo.broadcast_in_dim %1259, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1261 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1262 = stablehlo.broadcast_in_dim %1261, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1263 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1264 = stablehlo.broadcast_in_dim %1263, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1265 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1266 = stablehlo.broadcast_in_dim %1265, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1267 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1268 = stablehlo.broadcast_in_dim %1267, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1269 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1270 = stablehlo.broadcast_in_dim %1269, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1271 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1272 = stablehlo.broadcast_in_dim %1271, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1273 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1274 = stablehlo.broadcast_in_dim %1273, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1275 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1276 = stablehlo.broadcast_in_dim %1275, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1277 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1278 = stablehlo.broadcast_in_dim %1277, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1279 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1280 = stablehlo.broadcast_in_dim %1279, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1281 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1282 = stablehlo.broadcast_in_dim %1281, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1283 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1284 = stablehlo.broadcast_in_dim %1283, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1285 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1286 = stablehlo.broadcast_in_dim %1285, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1287 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1288 = stablehlo.broadcast_in_dim %1287, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1289 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1290 = stablehlo.broadcast_in_dim %1289, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1291 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1292 = stablehlo.broadcast_in_dim %1291, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1293 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1294 = stablehlo.broadcast_in_dim %1293, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1295 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1296 = stablehlo.broadcast_in_dim %1295, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1297 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1298 = stablehlo.broadcast_in_dim %1297, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1299 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1300 = stablehlo.broadcast_in_dim %1299, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1301 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1302 = stablehlo.broadcast_in_dim %1301, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1303 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1304 = stablehlo.broadcast_in_dim %1303, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1305 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1306 = stablehlo.broadcast_in_dim %1305, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1307 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1308 = stablehlo.broadcast_in_dim %1307, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1309 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1310 = stablehlo.broadcast_in_dim %1309, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1311 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1312 = stablehlo.broadcast_in_dim %1311, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1313 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1314 = stablehlo.broadcast_in_dim %1313, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1315 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1316 = stablehlo.broadcast_in_dim %1315, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1317 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1318 = stablehlo.broadcast_in_dim %1317, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1319 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1320 = stablehlo.broadcast_in_dim %1319, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1321 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1322 = stablehlo.broadcast_in_dim %1321, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1323 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1324 = stablehlo.broadcast_in_dim %1323, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1325 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1326 = stablehlo.broadcast_in_dim %1325, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1327 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1328 = stablehlo.broadcast_in_dim %1327, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1329 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1330 = stablehlo.broadcast_in_dim %1329, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1331 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1332 = stablehlo.broadcast_in_dim %1331, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1333 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1334 = stablehlo.broadcast_in_dim %1333, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1335 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1336 = stablehlo.broadcast_in_dim %1335, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1337 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1338 = stablehlo.broadcast_in_dim %1337, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1339 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1340 = stablehlo.broadcast_in_dim %1339, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1341 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1342 = stablehlo.broadcast_in_dim %1341, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1343 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1344 = stablehlo.broadcast_in_dim %1343, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1345 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1346 = stablehlo.broadcast_in_dim %1345, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1347 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1348 = stablehlo.broadcast_in_dim %1347, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1350 = stablehlo.broadcast_in_dim %1349, dims = [] : (tensor<f32>) -> tensor<1x1xf32> loc(#loc1961)
    %1351 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1352 = stablehlo.broadcast_in_dim %1351, dims = [] : (tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1962)
    %1353 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1354 = stablehlo.broadcast_in_dim %1353, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1355 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1356 = stablehlo.broadcast_in_dim %1355, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1357 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1358 = stablehlo.broadcast_in_dim %1357, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1359 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1360 = stablehlo.broadcast_in_dim %1359, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1361 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1362 = stablehlo.broadcast_in_dim %1361, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1363 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1364 = stablehlo.broadcast_in_dim %1363, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1365 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1366 = stablehlo.broadcast_in_dim %1365, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1367 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1368 = stablehlo.broadcast_in_dim %1367, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1369 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1370 = stablehlo.broadcast_in_dim %1369, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1371 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1372 = stablehlo.broadcast_in_dim %1371, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1373 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1374 = stablehlo.broadcast_in_dim %1373, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1375 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1376 = stablehlo.broadcast_in_dim %1375, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1377 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1378 = stablehlo.broadcast_in_dim %1377, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1379 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1380 = stablehlo.broadcast_in_dim %1379, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1381 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1382 = stablehlo.broadcast_in_dim %1381, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1383 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1384 = stablehlo.broadcast_in_dim %1383, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1385 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1386 = stablehlo.broadcast_in_dim %1385, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1387 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1388 = stablehlo.broadcast_in_dim %1387, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1389 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1390 = stablehlo.broadcast_in_dim %1389, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1391 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1392 = stablehlo.broadcast_in_dim %1391, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1393 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1394 = stablehlo.broadcast_in_dim %1393, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1395 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1396 = stablehlo.broadcast_in_dim %1395, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1397 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1398 = stablehlo.broadcast_in_dim %1397, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1399 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1400 = stablehlo.broadcast_in_dim %1399, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1401 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1402 = stablehlo.broadcast_in_dim %1401, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1403 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1404 = stablehlo.broadcast_in_dim %1403, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1405 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1406 = stablehlo.broadcast_in_dim %1405, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1407 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1408 = stablehlo.broadcast_in_dim %1407, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1409 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1410 = stablehlo.broadcast_in_dim %1409, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1411 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1412 = stablehlo.broadcast_in_dim %1411, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1413 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1414 = stablehlo.broadcast_in_dim %1413, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1415 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1416 = stablehlo.broadcast_in_dim %1415, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1417 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1418 = stablehlo.broadcast_in_dim %1417, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1419 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1420 = stablehlo.broadcast_in_dim %1419, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1421 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1422 = stablehlo.broadcast_in_dim %1421, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1423 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1424 = stablehlo.broadcast_in_dim %1423, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1425 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1426 = stablehlo.broadcast_in_dim %1425, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1427 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1428 = stablehlo.broadcast_in_dim %1427, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1429 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1430 = stablehlo.broadcast_in_dim %1429, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1431 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1432 = stablehlo.broadcast_in_dim %1431, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1433 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1434 = stablehlo.broadcast_in_dim %1433, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1435 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1436 = stablehlo.broadcast_in_dim %1435, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1437 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1438 = stablehlo.broadcast_in_dim %1437, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1439 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1440 = stablehlo.broadcast_in_dim %1439, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1441 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1442 = stablehlo.broadcast_in_dim %1441, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1443 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1444 = stablehlo.broadcast_in_dim %1443, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1445 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1446 = stablehlo.broadcast_in_dim %1445, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1447 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1448 = stablehlo.broadcast_in_dim %1447, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1449 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1450 = stablehlo.broadcast_in_dim %1449, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1451 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1452 = stablehlo.broadcast_in_dim %1451, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1453 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1454 = stablehlo.broadcast_in_dim %1453, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1455 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1456 = stablehlo.broadcast_in_dim %1455, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1457 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1458 = stablehlo.broadcast_in_dim %1457, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1459 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1460 = stablehlo.broadcast_in_dim %1459, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1461 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1462 = stablehlo.broadcast_in_dim %1461, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1463 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1464 = stablehlo.broadcast_in_dim %1463, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1465 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1466 = stablehlo.broadcast_in_dim %1465, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1467 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1468 = stablehlo.broadcast_in_dim %1467, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1469 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1470 = stablehlo.broadcast_in_dim %1469, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1471 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1472 = stablehlo.broadcast_in_dim %1471, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1473 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1474 = stablehlo.broadcast_in_dim %1473, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1475 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1476 = stablehlo.broadcast_in_dim %1475, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1477 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1478 = stablehlo.broadcast_in_dim %1477, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1479 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1480 = stablehlo.broadcast_in_dim %1479, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1481 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1482 = stablehlo.broadcast_in_dim %1481, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1483 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1484 = stablehlo.broadcast_in_dim %1483, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1485 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1486 = stablehlo.broadcast_in_dim %1485, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1487 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1488 = stablehlo.broadcast_in_dim %1487, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1489 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1490 = stablehlo.broadcast_in_dim %1489, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1491 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1492 = stablehlo.broadcast_in_dim %1491, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1493 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1494 = stablehlo.broadcast_in_dim %1493, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1495 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1496 = stablehlo.broadcast_in_dim %1495, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1497 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1498 = stablehlo.broadcast_in_dim %1497, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1499 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1500 = stablehlo.broadcast_in_dim %1499, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1501 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1502 = stablehlo.broadcast_in_dim %1501, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1503 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1504 = stablehlo.broadcast_in_dim %1503, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1505 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1506 = stablehlo.broadcast_in_dim %1505, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1507 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1508 = stablehlo.broadcast_in_dim %1507, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1509 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1510 = stablehlo.broadcast_in_dim %1509, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1511 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1512 = stablehlo.broadcast_in_dim %1511, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1513 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1514 = stablehlo.broadcast_in_dim %1513, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1515 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1516 = stablehlo.broadcast_in_dim %1515, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1517 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1518 = stablehlo.broadcast_in_dim %1517, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1519 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1520 = stablehlo.broadcast_in_dim %1519, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1521 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1522 = stablehlo.broadcast_in_dim %1521, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1523 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1524 = stablehlo.broadcast_in_dim %1523, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1525 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1526 = stablehlo.broadcast_in_dim %1525, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1527 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1528 = stablehlo.broadcast_in_dim %1527, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1529 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1530 = stablehlo.broadcast_in_dim %1529, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1531 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1532 = stablehlo.broadcast_in_dim %1531, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1533 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1534 = stablehlo.broadcast_in_dim %1533, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1535 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1536 = stablehlo.broadcast_in_dim %1535, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1537 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1538 = stablehlo.broadcast_in_dim %1537, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1539 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1540 = stablehlo.broadcast_in_dim %1539, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1541 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1542 = stablehlo.broadcast_in_dim %1541, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1543 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1544 = stablehlo.broadcast_in_dim %1543, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1545 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1546 = stablehlo.broadcast_in_dim %1545, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1547 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1548 = stablehlo.broadcast_in_dim %1547, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1549 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1550 = stablehlo.broadcast_in_dim %1549, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1551 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1552 = stablehlo.broadcast_in_dim %1551, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1553 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1554 = stablehlo.broadcast_in_dim %1553, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1555 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1556 = stablehlo.broadcast_in_dim %1555, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1557 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1558 = stablehlo.broadcast_in_dim %1557, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1559 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1560 = stablehlo.broadcast_in_dim %1559, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1561 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1562 = stablehlo.broadcast_in_dim %1561, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1563 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1564 = stablehlo.broadcast_in_dim %1563, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1565 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1566 = stablehlo.broadcast_in_dim %1565, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1567 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1568 = stablehlo.broadcast_in_dim %1567, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1569 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1570 = stablehlo.broadcast_in_dim %1569, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1571 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1572 = stablehlo.broadcast_in_dim %1571, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1573 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1574 = stablehlo.broadcast_in_dim %1573, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1575 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1576 = stablehlo.broadcast_in_dim %1575, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1577 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1578 = stablehlo.broadcast_in_dim %1577, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1579 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1580 = stablehlo.broadcast_in_dim %1579, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1581 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1582 = stablehlo.broadcast_in_dim %1581, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1583 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1584 = stablehlo.broadcast_in_dim %1583, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1585 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1586 = stablehlo.broadcast_in_dim %1585, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1587 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1588 = stablehlo.broadcast_in_dim %1587, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1589 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1590 = stablehlo.broadcast_in_dim %1589, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1591 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1592 = stablehlo.broadcast_in_dim %1591, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1593 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1594 = stablehlo.broadcast_in_dim %1593, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1595 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1596 = stablehlo.broadcast_in_dim %1595, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1597 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1598 = stablehlo.broadcast_in_dim %1597, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1599 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1600 = stablehlo.broadcast_in_dim %1599, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1601 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1602 = stablehlo.broadcast_in_dim %1601, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1603 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1604 = stablehlo.broadcast_in_dim %1603, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1605 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1606 = stablehlo.broadcast_in_dim %1605, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1607 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1608 = stablehlo.broadcast_in_dim %1607, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1609 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1610 = stablehlo.broadcast_in_dim %1609, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1611 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1612 = stablehlo.broadcast_in_dim %1611, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1613 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1614 = stablehlo.broadcast_in_dim %1613, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1615 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1616 = stablehlo.broadcast_in_dim %1615, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1617 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1618 = stablehlo.broadcast_in_dim %1617, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1619 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1620 = stablehlo.broadcast_in_dim %1619, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1621 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1622 = stablehlo.broadcast_in_dim %1621, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1623 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1624 = stablehlo.broadcast_in_dim %1623, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1625 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1626 = stablehlo.broadcast_in_dim %1625, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1627 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1628 = stablehlo.broadcast_in_dim %1627, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1629 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1630 = stablehlo.broadcast_in_dim %1629, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1631 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1632 = stablehlo.broadcast_in_dim %1631, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1633 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1634 = stablehlo.broadcast_in_dim %1633, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1635 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1636 = stablehlo.broadcast_in_dim %1635, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1637 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1638 = stablehlo.broadcast_in_dim %1637, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1639 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1640 = stablehlo.broadcast_in_dim %1639, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1641 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1642 = stablehlo.broadcast_in_dim %1641, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1643 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1644 = stablehlo.broadcast_in_dim %1643, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1645 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1646 = stablehlo.broadcast_in_dim %1645, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1647 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1648 = stablehlo.broadcast_in_dim %1647, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1649 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1650 = stablehlo.broadcast_in_dim %1649, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1651 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1652 = stablehlo.broadcast_in_dim %1651, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1653 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1654 = stablehlo.broadcast_in_dim %1653, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1655 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1656 = stablehlo.broadcast_in_dim %1655, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1657 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1658 = stablehlo.broadcast_in_dim %1657, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1659 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1660 = stablehlo.broadcast_in_dim %1659, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1661 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1662 = stablehlo.broadcast_in_dim %1661, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1663 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1664 = stablehlo.broadcast_in_dim %1663, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1665 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1666 = stablehlo.broadcast_in_dim %1665, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1667 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1668 = stablehlo.broadcast_in_dim %1667, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1669 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1670 = stablehlo.broadcast_in_dim %1669, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1671 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1672 = stablehlo.broadcast_in_dim %1671, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1673 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1674 = stablehlo.broadcast_in_dim %1673, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1675 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1676 = stablehlo.broadcast_in_dim %1675, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1677 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1678 = stablehlo.broadcast_in_dim %1677, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1679 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1680 = stablehlo.broadcast_in_dim %1679, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1681 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1682 = stablehlo.broadcast_in_dim %1681, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1683 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1684 = stablehlo.broadcast_in_dim %1683, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc1963)
    %1685 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1686 = stablehlo.broadcast_in_dim %1685, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc1964)
    %1687 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1688 = stablehlo.broadcast_in_dim %1687, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1689 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1690 = stablehlo.broadcast_in_dim %1689, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc1966)
    %1691 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1692 = stablehlo.broadcast_in_dim %1691, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1693 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1694 = stablehlo.broadcast_in_dim %1693, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1695 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1696 = stablehlo.broadcast_in_dim %1695, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1697 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1698 = stablehlo.broadcast_in_dim %1697, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc1965)
    %1699 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1700 = stablehlo.broadcast_in_dim %1699, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc1967)
    %1701 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1702 = stablehlo.broadcast_in_dim %1701, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc1968)
    %1703 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1704 = stablehlo.broadcast_in_dim %1703, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc1969)
    %1705 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1970)
    %1706:408 = stablehlo.while(%iterArg = %1252, %iterArg_0 = %1256, %iterArg_1 = %1258, %iterArg_2 = %1260, %iterArg_3 = %1262, %iterArg_4 = %1264, %iterArg_5 = %1266, %iterArg_6 = %1268, %iterArg_7 = %1270, %iterArg_8 = %1272, %iterArg_9 = %1274, %iterArg_10 = %1276, %iterArg_11 = %1278, %iterArg_12 = %1280, %iterArg_13 = %1282, %iterArg_14 = %1284, %iterArg_15 = %1286, %iterArg_16 = %1288, %iterArg_17 = %1290, %iterArg_18 = %1292, %iterArg_19 = %1294, %iterArg_20 = %1296, %iterArg_21 = %1298, %iterArg_22 = %1300, %iterArg_23 = %1302, %iterArg_24 = %1304, %iterArg_25 = %1306, %iterArg_26 = %1308, %iterArg_27 = %1310, %iterArg_28 = %1312, %iterArg_29 = %1314, %iterArg_30 = %1316, %iterArg_31 = %1318, %iterArg_32 = %1320, %iterArg_33 = %1322, %iterArg_34 = %1324, %iterArg_35 = %1326, %iterArg_36 = %1328, %iterArg_37 = %1330, %iterArg_38 = %1332, %iterArg_39 = %1334, %iterArg_40 = %1336, %iterArg_41 = %1338, %iterArg_42 = %1340, %iterArg_43 = %1342, %iterArg_44 = %1344, %iterArg_45 = %1346, %iterArg_46 = %1348, %iterArg_47 = %1350, %iterArg_48 = %676#541, %iterArg_49 = %676#542, %iterArg_50 = %256, %iterArg_51 = %250, %iterArg_52 = %253, %iterArg_53 = %17, %iterArg_54 = %18, %iterArg_55 = %19, %iterArg_56 = %20, %iterArg_57 = %21, %iterArg_58 = %22, %iterArg_59 = %23, %iterArg_60 = %24, %iterArg_61 = %25, %iterArg_62 = %26, %iterArg_63 = %27, %iterArg_64 = %28, %iterArg_65 = %29, %iterArg_66 = %30, %iterArg_67 = %31, %iterArg_68 = %32, %iterArg_69 = %33, %iterArg_70 = %34, %iterArg_71 = %35, %iterArg_72 = %36, %iterArg_73 = %37, %iterArg_74 = %38, %iterArg_75 = %39, %iterArg_76 = %40, %iterArg_77 = %41, %iterArg_78 = %42, %iterArg_79 = %43, %iterArg_80 = %44, %iterArg_81 = %45, %iterArg_82 = %46, %iterArg_83 = %47, %iterArg_84 = %48, %iterArg_85 = %49, %iterArg_86 = %50, %iterArg_87 = %51, %iterArg_88 = %52, %iterArg_89 = %53, %iterArg_90 = %54, %iterArg_91 = %55, %iterArg_92 = %56, %iterArg_93 = %57, %iterArg_94 = %58, %iterArg_95 = %59, %iterArg_96 = %60, %iterArg_97 = %61, %iterArg_98 = %62, %iterArg_99 = %63, %iterArg_100 = %64, %iterArg_101 = %65, %iterArg_102 = %66, %iterArg_103 = %67, %iterArg_104 = %68, %iterArg_105 = %69, %iterArg_106 = %70, %iterArg_107 = %71, %iterArg_108 = %72, %iterArg_109 = %73, %iterArg_110 = %74, %iterArg_111 = %75, %iterArg_112 = %76, %iterArg_113 = %77, %iterArg_114 = %78, %iterArg_115 = %79, %iterArg_116 = %80, %iterArg_117 = %81, %iterArg_118 = %82, %iterArg_119 = %83, %iterArg_120 = %84, %iterArg_121 = %85, %iterArg_122 = %86, %iterArg_123 = %87, %iterArg_124 = %88, %iterArg_125 = %89, %iterArg_126 = %90, %iterArg_127 = %91, %iterArg_128 = %92, %iterArg_129 = %93, %iterArg_130 = %94, %iterArg_131 = %95, %iterArg_132 = %97, %iterArg_133 = %98, %iterArg_134 = %99, %iterArg_135 = %100, %iterArg_136 = %101, %iterArg_137 = %102, %iterArg_138 = %103, %iterArg_139 = %104, %iterArg_140 = %105, %iterArg_141 = %106, %iterArg_142 = %107, %iterArg_143 = %108, %iterArg_144 = %109, %iterArg_145 = %110, %iterArg_146 = %111, %iterArg_147 = %112, %iterArg_148 = %113, %iterArg_149 = %114, %iterArg_150 = %115, %iterArg_151 = %116, %iterArg_152 = %117, %iterArg_153 = %118, %iterArg_154 = %119, %iterArg_155 = %120, %iterArg_156 = %121, %iterArg_157 = %122, %iterArg_158 = %123, %iterArg_159 = %124, %iterArg_160 = %125, %iterArg_161 = %126, %iterArg_162 = %127, %iterArg_163 = %128, %iterArg_164 = %129, %iterArg_165 = %130, %iterArg_166 = %131, %iterArg_167 = %132, %iterArg_168 = %133, %iterArg_169 = %134, %iterArg_170 = %135, %iterArg_171 = %136, %iterArg_172 = %137, %iterArg_173 = %138, %iterArg_174 = %139, %iterArg_175 = %140, %iterArg_176 = %141, %iterArg_177 = %142, %iterArg_178 = %143, %iterArg_179 = %144, %iterArg_180 = %145, %iterArg_181 = %146, %iterArg_182 = %147, %iterArg_183 = %148, %iterArg_184 = %149, %iterArg_185 = %150, %iterArg_186 = %151, %iterArg_187 = %152, %iterArg_188 = %153, %iterArg_189 = %154, %iterArg_190 = %155, %iterArg_191 = %156, %iterArg_192 = %157, %iterArg_193 = %158, %iterArg_194 = %159, %iterArg_195 = %160, %iterArg_196 = %161, %iterArg_197 = %162, %iterArg_198 = %163, %iterArg_199 = %164, %iterArg_200 = %165, %iterArg_201 = %166, %iterArg_202 = %167, %iterArg_203 = %168, %iterArg_204 = %169, %iterArg_205 = %170, %iterArg_206 = %171, %iterArg_207 = %172, %iterArg_208 = %173, %iterArg_209 = %174, %iterArg_210 = %175, %iterArg_211 = %176, %iterArg_212 = %177, %iterArg_213 = %178, %iterArg_214 = %179, %iterArg_215 = %180, %iterArg_216 = %181, %iterArg_217 = %182, %iterArg_218 = %183, %iterArg_219 = %184, %iterArg_220 = %185, %iterArg_221 = %186, %iterArg_222 = %187, %iterArg_223 = %188, %iterArg_224 = %189, %iterArg_225 = %190, %iterArg_226 = %191, %iterArg_227 = %192, %iterArg_228 = %1705, %iterArg_229 = %1352, %iterArg_230 = %1354, %iterArg_231 = %1356, %iterArg_232 = %1358, %iterArg_233 = %1360, %iterArg_234 = %1362, %iterArg_235 = %1364, %iterArg_236 = %1366, %iterArg_237 = %1368, %iterArg_238 = %1370, %iterArg_239 = %1372, %iterArg_240 = %1374, %iterArg_241 = %1376, %iterArg_242 = %1378, %iterArg_243 = %1380, %iterArg_244 = %1382, %iterArg_245 = %1384, %iterArg_246 = %1386, %iterArg_247 = %1388, %iterArg_248 = %1390, %iterArg_249 = %1392, %iterArg_250 = %1394, %iterArg_251 = %1396, %iterArg_252 = %1398, %iterArg_253 = %1400, %iterArg_254 = %1402, %iterArg_255 = %1404, %iterArg_256 = %1406, %iterArg_257 = %1408, %iterArg_258 = %1410, %iterArg_259 = %1412, %iterArg_260 = %1414, %iterArg_261 = %1416, %iterArg_262 = %1418, %iterArg_263 = %1420, %iterArg_264 = %1422, %iterArg_265 = %1424, %iterArg_266 = %1426, %iterArg_267 = %1428, %iterArg_268 = %1430, %iterArg_269 = %1432, %iterArg_270 = %1434, %iterArg_271 = %1436, %iterArg_272 = %1438, %iterArg_273 = %1440, %iterArg_274 = %1442, %iterArg_275 = %1444, %iterArg_276 = %1446, %iterArg_277 = %1448, %iterArg_278 = %1450, %iterArg_279 = %1452, %iterArg_280 = %1454, %iterArg_281 = %1456, %iterArg_282 = %1458, %iterArg_283 = %1460, %iterArg_284 = %1462, %iterArg_285 = %1464, %iterArg_286 = %1466, %iterArg_287 = %1468, %iterArg_288 = %1470, %iterArg_289 = %1472, %iterArg_290 = %1474, %iterArg_291 = %1476, %iterArg_292 = %1478, %iterArg_293 = %1480, %iterArg_294 = %1482, %iterArg_295 = %1484, %iterArg_296 = %1486, %iterArg_297 = %1488, %iterArg_298 = %1490, %iterArg_299 = %1492, %iterArg_300 = %1494, %iterArg_301 = %1496, %iterArg_302 = %1498, %iterArg_303 = %1500, %iterArg_304 = %1502, %iterArg_305 = %1504, %iterArg_306 = %1506, %iterArg_307 = %1508, %iterArg_308 = %1510, %iterArg_309 = %1512, %iterArg_310 = %1514, %iterArg_311 = %1516, %iterArg_312 = %1518, %iterArg_313 = %1520, %iterArg_314 = %1522, %iterArg_315 = %1524, %iterArg_316 = %1526, %iterArg_317 = %1528, %iterArg_318 = %1530, %iterArg_319 = %1532, %iterArg_320 = %1534, %iterArg_321 = %1536, %iterArg_322 = %1538, %iterArg_323 = %1540, %iterArg_324 = %1542, %iterArg_325 = %1544, %iterArg_326 = %1546, %iterArg_327 = %1548, %iterArg_328 = %1550, %iterArg_329 = %1552, %iterArg_330 = %1554, %iterArg_331 = %1556, %iterArg_332 = %1558, %iterArg_333 = %1560, %iterArg_334 = %1562, %iterArg_335 = %1564, %iterArg_336 = %1566, %iterArg_337 = %1568, %iterArg_338 = %1570, %iterArg_339 = %1572, %iterArg_340 = %1574, %iterArg_341 = %1576, %iterArg_342 = %1578, %iterArg_343 = %1580, %iterArg_344 = %1582, %iterArg_345 = %1584, %iterArg_346 = %1586, %iterArg_347 = %1588, %iterArg_348 = %1590, %iterArg_349 = %1592, %iterArg_350 = %1594, %iterArg_351 = %1596, %iterArg_352 = %1598, %iterArg_353 = %1600, %iterArg_354 = %1602, %iterArg_355 = %1604, %iterArg_356 = %1606, %iterArg_357 = %1608, %iterArg_358 = %1610, %iterArg_359 = %1612, %iterArg_360 = %1614, %iterArg_361 = %1616, %iterArg_362 = %1618, %iterArg_363 = %1620, %iterArg_364 = %1622, %iterArg_365 = %1624, %iterArg_366 = %1626, %iterArg_367 = %1628, %iterArg_368 = %1630, %iterArg_369 = %1632, %iterArg_370 = %1634, %iterArg_371 = %1636, %iterArg_372 = %1638, %iterArg_373 = %1640, %iterArg_374 = %1642, %iterArg_375 = %1644, %iterArg_376 = %1646, %iterArg_377 = %1648, %iterArg_378 = %1650, %iterArg_379 = %1652, %iterArg_380 = %1654, %iterArg_381 = %1656, %iterArg_382 = %1658, %iterArg_383 = %1660, %iterArg_384 = %1662, %iterArg_385 = %1664, %iterArg_386 = %1666, %iterArg_387 = %1668, %iterArg_388 = %1670, %iterArg_389 = %1672, %iterArg_390 = %1674, %iterArg_391 = %1676, %iterArg_392 = %1678, %iterArg_393 = %1680, %iterArg_394 = %1682, %iterArg_395 = %1684, %iterArg_396 = %1686, %iterArg_397 = %1688, %iterArg_398 = %1690, %iterArg_399 = %1692, %iterArg_400 = %1694, %iterArg_401 = %1696, %iterArg_402 = %1698, %iterArg_403 = %1700, %iterArg_404 = %1702, %iterArg_405 = %1704, %iterArg_406 = %1254) : tensor<1x1x1x2048x2048xbf16>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1x1x2048x2048xbf16>, tensor<1xi32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1x2048xi32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<i32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>
     cond {
      %29929 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29930 = stablehlo.compare  LT, %iterArg_228, %29929,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1972)
      stablehlo.return %29930 : tensor<i1> loc(#loc1971)
    } do {
      %29929 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29930 = stablehlo.subtract %29929, %iterArg_228 : tensor<i32> loc(#loc1973)
      %29931 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29932 = stablehlo.subtract %29930, %29931 : tensor<i32> loc(#loc1973)
      %29933 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29934 = stablehlo.compare  LT, %29932, %29933,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29935 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29936 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29937 = stablehlo.add %29935, %29936 : tensor<i32> loc(#loc1976)
      %29938 = stablehlo.select %29934, %29937, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29939 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29940 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29941 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29942 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29943 = stablehlo.dynamic_slice %iterArg, %29938, %29939, %29940, %29941, %29942, sizes = [1, 1, 1, 2048, 2048] : (tensor<1x1x1x2048x2048xbf16>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1978)
      %29944 = stablehlo.reshape %29943 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1979)
      %29945 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29946 = stablehlo.compare  LT, %29932, %29945,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29947 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29948 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29949 = stablehlo.add %29947, %29948 : tensor<i32> loc(#loc1976)
      %29950 = stablehlo.select %29946, %29949, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29951 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29952 = stablehlo.dynamic_slice %iterArg_0, %29950, %29951, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %29953 = stablehlo.reshape %29952 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %29954 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29955 = stablehlo.compare  LT, %29932, %29954,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29956 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29957 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29958 = stablehlo.add %29956, %29957 : tensor<i32> loc(#loc1976)
      %29959 = stablehlo.select %29955, %29958, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29960 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29961 = stablehlo.dynamic_slice %iterArg_1, %29959, %29960, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %29962 = stablehlo.reshape %29961 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %29963 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29964 = stablehlo.compare  LT, %29932, %29963,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29965 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29966 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29967 = stablehlo.add %29965, %29966 : tensor<i32> loc(#loc1976)
      %29968 = stablehlo.select %29964, %29967, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29969 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29970 = stablehlo.dynamic_slice %iterArg_2, %29968, %29969, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %29971 = stablehlo.reshape %29970 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %29972 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29973 = stablehlo.compare  LT, %29932, %29972,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29974 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29975 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29976 = stablehlo.add %29974, %29975 : tensor<i32> loc(#loc1976)
      %29977 = stablehlo.select %29973, %29976, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29978 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29979 = stablehlo.dynamic_slice %iterArg_3, %29977, %29978, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %29980 = stablehlo.reshape %29979 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %29981 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29982 = stablehlo.compare  LT, %29932, %29981,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29983 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29984 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29985 = stablehlo.add %29983, %29984 : tensor<i32> loc(#loc1976)
      %29986 = stablehlo.select %29982, %29985, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29987 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29988 = stablehlo.dynamic_slice %iterArg_4, %29986, %29987, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %29989 = stablehlo.reshape %29988 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %29990 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29991 = stablehlo.compare  LT, %29932, %29990,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %29992 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %29993 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %29994 = stablehlo.add %29992, %29993 : tensor<i32> loc(#loc1976)
      %29995 = stablehlo.select %29991, %29994, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %29996 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %29997 = stablehlo.dynamic_slice %iterArg_5, %29995, %29996, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %29998 = stablehlo.reshape %29997 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %29999 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30000 = stablehlo.compare  LT, %29932, %29999,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30001 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30002 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30003 = stablehlo.add %30001, %30002 : tensor<i32> loc(#loc1976)
      %30004 = stablehlo.select %30000, %30003, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30005 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30006 = stablehlo.dynamic_slice %iterArg_6, %30004, %30005, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30007 = stablehlo.reshape %30006 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30008 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30009 = stablehlo.compare  LT, %29932, %30008,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30010 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30011 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30012 = stablehlo.add %30010, %30011 : tensor<i32> loc(#loc1976)
      %30013 = stablehlo.select %30009, %30012, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30014 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30015 = stablehlo.dynamic_slice %iterArg_7, %30013, %30014, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30016 = stablehlo.reshape %30015 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30017 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30018 = stablehlo.compare  LT, %29932, %30017,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30019 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30020 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30021 = stablehlo.add %30019, %30020 : tensor<i32> loc(#loc1976)
      %30022 = stablehlo.select %30018, %30021, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30023 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30024 = stablehlo.dynamic_slice %iterArg_8, %30022, %30023, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30025 = stablehlo.reshape %30024 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30026 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30027 = stablehlo.compare  LT, %29932, %30026,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30028 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30029 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30030 = stablehlo.add %30028, %30029 : tensor<i32> loc(#loc1976)
      %30031 = stablehlo.select %30027, %30030, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30032 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30033 = stablehlo.dynamic_slice %iterArg_9, %30031, %30032, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30034 = stablehlo.reshape %30033 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30035 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30036 = stablehlo.compare  LT, %29932, %30035,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30037 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30038 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30039 = stablehlo.add %30037, %30038 : tensor<i32> loc(#loc1976)
      %30040 = stablehlo.select %30036, %30039, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30041 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30042 = stablehlo.dynamic_slice %iterArg_10, %30040, %30041, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30043 = stablehlo.reshape %30042 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30044 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30045 = stablehlo.compare  LT, %29932, %30044,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30046 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30047 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30048 = stablehlo.add %30046, %30047 : tensor<i32> loc(#loc1976)
      %30049 = stablehlo.select %30045, %30048, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30050 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30051 = stablehlo.dynamic_slice %iterArg_11, %30049, %30050, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30052 = stablehlo.reshape %30051 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30053 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30054 = stablehlo.compare  LT, %29932, %30053,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30055 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30056 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30057 = stablehlo.add %30055, %30056 : tensor<i32> loc(#loc1976)
      %30058 = stablehlo.select %30054, %30057, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30059 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30060 = stablehlo.dynamic_slice %iterArg_12, %30058, %30059, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30061 = stablehlo.reshape %30060 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30062 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30063 = stablehlo.compare  LT, %29932, %30062,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30064 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30065 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30066 = stablehlo.add %30064, %30065 : tensor<i32> loc(#loc1976)
      %30067 = stablehlo.select %30063, %30066, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30068 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30069 = stablehlo.dynamic_slice %iterArg_13, %30067, %30068, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30070 = stablehlo.reshape %30069 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30071 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30072 = stablehlo.compare  LT, %29932, %30071,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30073 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30074 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30075 = stablehlo.add %30073, %30074 : tensor<i32> loc(#loc1976)
      %30076 = stablehlo.select %30072, %30075, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30077 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30078 = stablehlo.dynamic_slice %iterArg_14, %30076, %30077, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30079 = stablehlo.reshape %30078 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30080 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30081 = stablehlo.compare  LT, %29932, %30080,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30082 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30083 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30084 = stablehlo.add %30082, %30083 : tensor<i32> loc(#loc1976)
      %30085 = stablehlo.select %30081, %30084, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30086 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30087 = stablehlo.dynamic_slice %iterArg_15, %30085, %30086, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30088 = stablehlo.reshape %30087 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30089 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30090 = stablehlo.compare  LT, %29932, %30089,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30091 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30092 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30093 = stablehlo.add %30091, %30092 : tensor<i32> loc(#loc1976)
      %30094 = stablehlo.select %30090, %30093, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30095 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30096 = stablehlo.dynamic_slice %iterArg_16, %30094, %30095, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30097 = stablehlo.reshape %30096 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30098 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30099 = stablehlo.compare  LT, %29932, %30098,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30100 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30101 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30102 = stablehlo.add %30100, %30101 : tensor<i32> loc(#loc1976)
      %30103 = stablehlo.select %30099, %30102, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30104 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30105 = stablehlo.dynamic_slice %iterArg_17, %30103, %30104, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30106 = stablehlo.reshape %30105 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30107 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30108 = stablehlo.compare  LT, %29932, %30107,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30109 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30110 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30111 = stablehlo.add %30109, %30110 : tensor<i32> loc(#loc1976)
      %30112 = stablehlo.select %30108, %30111, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30113 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30114 = stablehlo.dynamic_slice %iterArg_18, %30112, %30113, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30115 = stablehlo.reshape %30114 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30116 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30117 = stablehlo.compare  LT, %29932, %30116,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30118 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30119 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30120 = stablehlo.add %30118, %30119 : tensor<i32> loc(#loc1976)
      %30121 = stablehlo.select %30117, %30120, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30122 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30123 = stablehlo.dynamic_slice %iterArg_19, %30121, %30122, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30124 = stablehlo.reshape %30123 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30125 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30126 = stablehlo.compare  LT, %29932, %30125,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30127 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30128 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30129 = stablehlo.add %30127, %30128 : tensor<i32> loc(#loc1976)
      %30130 = stablehlo.select %30126, %30129, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30131 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30132 = stablehlo.dynamic_slice %iterArg_20, %30130, %30131, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30133 = stablehlo.reshape %30132 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30134 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30135 = stablehlo.compare  LT, %29932, %30134,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30136 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30137 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30138 = stablehlo.add %30136, %30137 : tensor<i32> loc(#loc1976)
      %30139 = stablehlo.select %30135, %30138, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30140 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30141 = stablehlo.dynamic_slice %iterArg_21, %30139, %30140, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30142 = stablehlo.reshape %30141 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30143 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30144 = stablehlo.compare  LT, %29932, %30143,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30145 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30146 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30147 = stablehlo.add %30145, %30146 : tensor<i32> loc(#loc1976)
      %30148 = stablehlo.select %30144, %30147, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30149 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30150 = stablehlo.dynamic_slice %iterArg_22, %30148, %30149, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30151 = stablehlo.reshape %30150 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30152 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30153 = stablehlo.compare  LT, %29932, %30152,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30154 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30155 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30156 = stablehlo.add %30154, %30155 : tensor<i32> loc(#loc1976)
      %30157 = stablehlo.select %30153, %30156, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30158 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30159 = stablehlo.dynamic_slice %iterArg_23, %30157, %30158, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30160 = stablehlo.reshape %30159 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30161 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30162 = stablehlo.compare  LT, %29932, %30161,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30163 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30164 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30165 = stablehlo.add %30163, %30164 : tensor<i32> loc(#loc1976)
      %30166 = stablehlo.select %30162, %30165, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30167 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30168 = stablehlo.dynamic_slice %iterArg_24, %30166, %30167, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30169 = stablehlo.reshape %30168 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30170 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30171 = stablehlo.compare  LT, %29932, %30170,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30172 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30173 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30174 = stablehlo.add %30172, %30173 : tensor<i32> loc(#loc1976)
      %30175 = stablehlo.select %30171, %30174, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30176 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30177 = stablehlo.dynamic_slice %iterArg_25, %30175, %30176, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30178 = stablehlo.reshape %30177 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30179 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30180 = stablehlo.compare  LT, %29932, %30179,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30181 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30182 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30183 = stablehlo.add %30181, %30182 : tensor<i32> loc(#loc1976)
      %30184 = stablehlo.select %30180, %30183, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30185 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30186 = stablehlo.dynamic_slice %iterArg_26, %30184, %30185, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30187 = stablehlo.reshape %30186 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30188 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30189 = stablehlo.compare  LT, %29932, %30188,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30190 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30191 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30192 = stablehlo.add %30190, %30191 : tensor<i32> loc(#loc1976)
      %30193 = stablehlo.select %30189, %30192, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30194 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30195 = stablehlo.dynamic_slice %iterArg_27, %30193, %30194, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30196 = stablehlo.reshape %30195 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30197 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30198 = stablehlo.compare  LT, %29932, %30197,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30199 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30200 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30201 = stablehlo.add %30199, %30200 : tensor<i32> loc(#loc1976)
      %30202 = stablehlo.select %30198, %30201, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30203 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30204 = stablehlo.dynamic_slice %iterArg_28, %30202, %30203, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30205 = stablehlo.reshape %30204 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30206 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30207 = stablehlo.compare  LT, %29932, %30206,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30208 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30209 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30210 = stablehlo.add %30208, %30209 : tensor<i32> loc(#loc1976)
      %30211 = stablehlo.select %30207, %30210, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30212 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30213 = stablehlo.dynamic_slice %iterArg_29, %30211, %30212, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30214 = stablehlo.reshape %30213 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30215 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30216 = stablehlo.compare  LT, %29932, %30215,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30217 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30218 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30219 = stablehlo.add %30217, %30218 : tensor<i32> loc(#loc1976)
      %30220 = stablehlo.select %30216, %30219, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30221 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30222 = stablehlo.dynamic_slice %iterArg_30, %30220, %30221, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30223 = stablehlo.reshape %30222 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30224 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30225 = stablehlo.compare  LT, %29932, %30224,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30226 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30227 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30228 = stablehlo.add %30226, %30227 : tensor<i32> loc(#loc1976)
      %30229 = stablehlo.select %30225, %30228, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30230 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30231 = stablehlo.dynamic_slice %iterArg_31, %30229, %30230, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30232 = stablehlo.reshape %30231 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30233 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30234 = stablehlo.compare  LT, %29932, %30233,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30235 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30236 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30237 = stablehlo.add %30235, %30236 : tensor<i32> loc(#loc1976)
      %30238 = stablehlo.select %30234, %30237, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30239 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30240 = stablehlo.dynamic_slice %iterArg_32, %30238, %30239, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30241 = stablehlo.reshape %30240 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30242 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30243 = stablehlo.compare  LT, %29932, %30242,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30244 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30245 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30246 = stablehlo.add %30244, %30245 : tensor<i32> loc(#loc1976)
      %30247 = stablehlo.select %30243, %30246, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30248 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30249 = stablehlo.dynamic_slice %iterArg_33, %30247, %30248, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30250 = stablehlo.reshape %30249 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30251 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30252 = stablehlo.compare  LT, %29932, %30251,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30253 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30254 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30255 = stablehlo.add %30253, %30254 : tensor<i32> loc(#loc1976)
      %30256 = stablehlo.select %30252, %30255, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30257 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30258 = stablehlo.dynamic_slice %iterArg_34, %30256, %30257, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30259 = stablehlo.reshape %30258 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30260 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30261 = stablehlo.compare  LT, %29932, %30260,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30262 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30263 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30264 = stablehlo.add %30262, %30263 : tensor<i32> loc(#loc1976)
      %30265 = stablehlo.select %30261, %30264, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30266 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30267 = stablehlo.dynamic_slice %iterArg_35, %30265, %30266, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30268 = stablehlo.reshape %30267 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30269 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30270 = stablehlo.compare  LT, %29932, %30269,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30271 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30272 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30273 = stablehlo.add %30271, %30272 : tensor<i32> loc(#loc1976)
      %30274 = stablehlo.select %30270, %30273, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30275 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30276 = stablehlo.dynamic_slice %iterArg_36, %30274, %30275, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30277 = stablehlo.reshape %30276 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30278 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30279 = stablehlo.compare  LT, %29932, %30278,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30280 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30281 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30282 = stablehlo.add %30280, %30281 : tensor<i32> loc(#loc1976)
      %30283 = stablehlo.select %30279, %30282, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30284 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30285 = stablehlo.dynamic_slice %iterArg_37, %30283, %30284, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30286 = stablehlo.reshape %30285 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30287 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30288 = stablehlo.compare  LT, %29932, %30287,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30289 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30290 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30291 = stablehlo.add %30289, %30290 : tensor<i32> loc(#loc1976)
      %30292 = stablehlo.select %30288, %30291, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30293 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30294 = stablehlo.dynamic_slice %iterArg_38, %30292, %30293, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30295 = stablehlo.reshape %30294 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30296 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30297 = stablehlo.compare  LT, %29932, %30296,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30298 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30299 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30300 = stablehlo.add %30298, %30299 : tensor<i32> loc(#loc1976)
      %30301 = stablehlo.select %30297, %30300, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30302 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30303 = stablehlo.dynamic_slice %iterArg_39, %30301, %30302, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30304 = stablehlo.reshape %30303 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30305 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30306 = stablehlo.compare  LT, %29932, %30305,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30307 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30308 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30309 = stablehlo.add %30307, %30308 : tensor<i32> loc(#loc1976)
      %30310 = stablehlo.select %30306, %30309, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30311 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30312 = stablehlo.dynamic_slice %iterArg_40, %30310, %30311, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30313 = stablehlo.reshape %30312 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30314 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30315 = stablehlo.compare  LT, %29932, %30314,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30316 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30317 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30318 = stablehlo.add %30316, %30317 : tensor<i32> loc(#loc1976)
      %30319 = stablehlo.select %30315, %30318, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30320 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30321 = stablehlo.dynamic_slice %iterArg_41, %30319, %30320, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30322 = stablehlo.reshape %30321 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30323 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30324 = stablehlo.compare  LT, %29932, %30323,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30325 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30326 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30327 = stablehlo.add %30325, %30326 : tensor<i32> loc(#loc1976)
      %30328 = stablehlo.select %30324, %30327, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30329 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30330 = stablehlo.dynamic_slice %iterArg_42, %30328, %30329, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30331 = stablehlo.reshape %30330 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30332 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30333 = stablehlo.compare  LT, %29932, %30332,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30334 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30335 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30336 = stablehlo.add %30334, %30335 : tensor<i32> loc(#loc1976)
      %30337 = stablehlo.select %30333, %30336, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30338 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30339 = stablehlo.dynamic_slice %iterArg_43, %30337, %30338, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30340 = stablehlo.reshape %30339 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30341 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30342 = stablehlo.compare  LT, %29932, %30341,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30343 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30344 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30345 = stablehlo.add %30343, %30344 : tensor<i32> loc(#loc1976)
      %30346 = stablehlo.select %30342, %30345, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30347 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30348 = stablehlo.dynamic_slice %iterArg_44, %30346, %30347, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30349 = stablehlo.reshape %30348 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30350 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30351 = stablehlo.compare  LT, %29932, %30350,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30352 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30353 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30354 = stablehlo.add %30352, %30353 : tensor<i32> loc(#loc1976)
      %30355 = stablehlo.select %30351, %30354, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30356 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30357 = stablehlo.dynamic_slice %iterArg_45, %30355, %30356, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30358 = stablehlo.reshape %30357 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30359 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30360 = stablehlo.compare  LT, %29932, %30359,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30361 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30362 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30363 = stablehlo.add %30361, %30362 : tensor<i32> loc(#loc1976)
      %30364 = stablehlo.select %30360, %30363, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30365 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30366 = stablehlo.dynamic_slice %iterArg_46, %30364, %30365, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30367 = stablehlo.reshape %30366 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30368 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30369 = stablehlo.compare  LT, %29932, %30368,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30370 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30371 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30372 = stablehlo.add %30370, %30371 : tensor<i32> loc(#loc1976)
      %30373 = stablehlo.select %30369, %30372, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30374 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30375 = stablehlo.dynamic_slice %iterArg_47, %30373, %30374, sizes = [1, 1] : (tensor<1x1xf32>, tensor<i32>, tensor<i32>) -> tensor<1x1xf32> loc(#loc1980)
      %30376 = stablehlo.reshape %30375 : (tensor<1x1xf32>) -> tensor<1xf32> loc(#loc1979)
      %30377 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30378 = stablehlo.compare  LT, %29932, %30377,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30379 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30380 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30381 = stablehlo.add %30379, %30380 : tensor<i32> loc(#loc1976)
      %30382 = stablehlo.select %30378, %30381, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30383 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30384 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30385 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30386 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30387 = stablehlo.dynamic_slice %iterArg_48, %30382, %30383, %30384, %30385, %30386, sizes = [1, 1, 1, 2048, 2048] : (tensor<1x1x1x2048x2048xbf16>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc1978)
      %30388 = stablehlo.reshape %30387 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1979)
      %30389 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1971)
      %30390 = stablehlo.compare  LT, %29932, %30389,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1974)
      %30391 = stablehlo.convert %29932 : tensor<i32> loc(#loc1975)
      %30392 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %30393 = stablehlo.add %30391, %30392 : tensor<i32> loc(#loc1976)
      %30394 = stablehlo.select %30390, %30393, %29932 : tensor<i1>, tensor<i32> loc(#loc1977)
      %30395 = stablehlo.dynamic_slice %iterArg_49, %30394, sizes = [1] : (tensor<1xi32>, tensor<i32>) -> tensor<1xi32> loc(#loc1981)
      %30396 = stablehlo.reshape %30395 : (tensor<1xi32>) -> tensor<i32> loc(#loc1979)
      %30397 = stablehlo.iota dim = 0 : tensor<1xi32> loc(#loc1982)
      %30398 = stablehlo.broadcast_in_dim %30396, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1983)
      %30399 = stablehlo.subtract %30398, %30397 : tensor<1xi32> loc(#loc1983)
      %30400 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30401 = stablehlo.broadcast_in_dim %30400, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1985)
      %30402 = stablehlo.maximum %30399, %30401 : tensor<1xi32> loc(#loc1985)
      %30403 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30404 = func.call @remainder_459(%30402, %30403) : (tensor<1xi32>, tensor<i32>) -> tensor<1xi32> loc(#loc1986)
      %30405 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30406 = func.call @remainder_461(%30396, %30405) : (tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc1987)
      %30407 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30408 = stablehlo.compare  LT, %30406, %30407,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc1988)
      %30409 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30410 = stablehlo.add %30406, %30409 : tensor<i32> loc(#loc1989)
      %30411 = stablehlo.select %30408, %30410, %30406 : tensor<i1>, tensor<i32> loc(#loc1990)
      %30412 = stablehlo.broadcast_in_dim %30411, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1991)
      %30413 = "stablehlo.gather"(%iterArg_50, %30412) {dimension_numbers = #stablehlo.gather<offset_dims = [0, 1, 2, 3], collapsed_slice_dims = [0], start_index_map = [0]>, indices_are_sorted = true, slice_sizes = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>, tensor<1xi32>) -> tensor<1x1x2048x2048xbf16> loc(#loc1992)
      %30414 = stablehlo.iota dim = 0 : tensor<1x1x2048x2048xi32> loc(#loc1993)
      %30415 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30416 = stablehlo.broadcast_in_dim %30415, dims = [] : (tensor<i32>) -> tensor<1x1x2048x2048xi32> loc(#loc1994)
      %30417 = stablehlo.compare  EQ, %30414, %30416,  SIGNED : (tensor<1x1x2048x2048xi32>, tensor<1x1x2048x2048xi32>) -> tensor<1x1x2048x2048xi1> loc(#loc1994)
      %30418 = func.call @_where_463(%30417, %30413, %30388) : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc1995)
      %30419 = arith.select %ptrue, %30404, %30404 : tensor<1xi32>
      %30420 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30421 = stablehlo.broadcast_in_dim %30420, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1997)
      %30422 = stablehlo.compare  LT, %30419, %30421,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc1997)
      %30423 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30424 = stablehlo.broadcast_in_dim %30423, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1998)
      %30425 = stablehlo.add %30419, %30424 : tensor<1xi32> loc(#loc1998)
      %30426 = stablehlo.select %30422, %30425, %30419 : tensor<1xi1>, tensor<1xi32> loc(#loc1999)
      %30427 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30428 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30429 = stablehlo.add %30427, %30428 : tensor<i32> loc(#loc1998)
      %30430 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %30431 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30432 = stablehlo.select %30430, %30429, %30431 : tensor<i1>, tensor<i32> loc(#loc1999)
      %30433 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30434 = stablehlo.constant dense<2048> : tensor<i32> loc(#loc1984)
      %30435 = stablehlo.add %30433, %30434 : tensor<i32> loc(#loc1998)
      %30436 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %30437 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30438 = stablehlo.select %30436, %30435, %30437 : tensor<i1>, tensor<i32> loc(#loc1999)
      %30439 = stablehlo.broadcast_in_dim %30426, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32> loc(#loc2000)
      %30440 = stablehlo.broadcast_in_dim %30432, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc2001)
      %30441 = stablehlo.broadcast_in_dim %30438, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc2001)
      %30442 = stablehlo.concatenate %30439, %30440, %30441, dim = 1 : (tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<1x3xi32> loc(#loc2002)
      %30443 = "stablehlo.gather"(%iterArg_51, %30442) {dimension_numbers = #stablehlo.gather<offset_dims = [1, 2, 3], start_index_map = [0, 1, 2], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = dense<[1, 1, 2048]> : tensor<3xi64>} : (tensor<1x1x2048xi32>, tensor<1x3xi32>) -> tensor<1x1x1x2048xi32> loc(#loc2003)
      %30444 = stablehlo.reshape %30443 : (tensor<1x1x1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc2004)
      %30445 = arith.select %ptrue, %30444, %30444 : tensor<1x1x2048xi32>
      %30446 = arith.select %ptrue, %30404, %30404 : tensor<1xi32>
      %30447 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30448 = stablehlo.broadcast_in_dim %30447, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1997)
      %30449 = stablehlo.compare  LT, %30446, %30448,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc1997)
      %30450 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30451 = stablehlo.broadcast_in_dim %30450, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc1998)
      %30452 = stablehlo.add %30446, %30451 : tensor<1xi32> loc(#loc1998)
      %30453 = stablehlo.select %30449, %30452, %30446 : tensor<1xi1>, tensor<1xi32> loc(#loc1999)
      %30454 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30455 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30456 = stablehlo.add %30454, %30455 : tensor<i32> loc(#loc1998)
      %30457 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %30458 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30459 = stablehlo.select %30457, %30456, %30458 : tensor<i1>, tensor<i32> loc(#loc1999)
      %30460 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30461 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30462 = stablehlo.add %30460, %30461 : tensor<i32> loc(#loc1998)
      %30463 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %30464 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30465 = stablehlo.select %30463, %30462, %30464 : tensor<i1>, tensor<i32> loc(#loc1999)
      %30466 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30467 = stablehlo.constant dense<2048> : tensor<i32> loc(#loc1984)
      %30468 = stablehlo.add %30466, %30467 : tensor<i32> loc(#loc1998)
      %30469 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %30470 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30471 = stablehlo.select %30469, %30468, %30470 : tensor<i1>, tensor<i32> loc(#loc1999)
      %30472 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30473 = stablehlo.constant dense<2048> : tensor<i32> loc(#loc1984)
      %30474 = stablehlo.add %30472, %30473 : tensor<i32> loc(#loc1998)
      %30475 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %30476 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1984)
      %30477 = stablehlo.select %30475, %30474, %30476 : tensor<i1>, tensor<i32> loc(#loc1999)
      %30478 = stablehlo.broadcast_in_dim %30453, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32> loc(#loc2000)
      %30479 = stablehlo.broadcast_in_dim %30459, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc2001)
      %30480 = stablehlo.broadcast_in_dim %30465, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc2001)
      %30481 = stablehlo.broadcast_in_dim %30471, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc2001)
      %30482 = stablehlo.broadcast_in_dim %30477, dims = [] : (tensor<i32>) -> tensor<1x1xi32> loc(#loc2001)
      %30483 = stablehlo.concatenate %30478, %30479, %30480, %30481, %30482, dim = 1 : (tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<1x5xi32> loc(#loc2002)
      %30484 = "stablehlo.gather"(%iterArg_52, %30483) {dimension_numbers = #stablehlo.gather<offset_dims = [1, 2, 3, 4, 5], start_index_map = [0, 1, 2, 3, 4], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>, tensor<1x5xi32>) -> tensor<1x1x1x1x2048x2048xbf16> loc(#loc2006)
      %30485 = stablehlo.reshape %30484 : (tensor<1x1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc2004)
      %30486 = arith.select %ptrue, %30485, %30485 : tensor<1x1x1x2048x2048xbf16>
      %30487 = stablehlo.iota dim = 0 : tensor<1xi32> loc(#loc2008)
      %30488 = stablehlo.broadcast_in_dim %30396, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc2009)
      %30489 = stablehlo.compare  LE, %30487, %30488,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc2009)
      %30490 = stablehlo.broadcast_in_dim %30396, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc2010)
      %30491 = stablehlo.subtract %30490, %30487 : tensor<1xi32> loc(#loc2010)
      %30492 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1984)
      %30493 = stablehlo.broadcast_in_dim %30492, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc2011)
      %30494 = stablehlo.compare  LT, %30491, %30493,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc2011)
      %30495 = stablehlo.and %30489, %30494 : tensor<1xi1> loc(#loc2012)
      %30496 = arith.select %ptrue, %30418, %30418 : tensor<1x1x2048x2048xbf16>
      %30497 = arith.select %ptrue, %30445, %30445 : tensor<1x1x2048xi32>
      %30498 = arith.select %ptrue, %30486, %30486 : tensor<1x1x1x2048x2048xbf16>
      %30499:2 = func.call @_where_464(%30495, %iterArg_53, %iterArg_53) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30500:2 = func.call @_where_465(%30495, %iterArg_54, %iterArg_54) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30501:2 = func.call @_where_466(%30495, %iterArg_55, %iterArg_55) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30502:2 = func.call @_where_467(%30495, %iterArg_56, %iterArg_56) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30503:2 = func.call @_where_468(%30495, %iterArg_57, %iterArg_57) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30504:2 = func.call @_where_469(%30495, %iterArg_58, %iterArg_58) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30505:2 = func.call @_where_470(%30495, %iterArg_59, %iterArg_59) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30506:2 = func.call @_where_471(%30495, %iterArg_60, %iterArg_60) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30507:2 = func.call @_where_472(%30495, %iterArg_61, %iterArg_61) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30508:2 = func.call @_where_473(%30495, %iterArg_62, %iterArg_62) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30509:2 = func.call @_where_474(%30495, %iterArg_63, %iterArg_63) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30510:2 = func.call @_where_475(%30495, %iterArg_64, %iterArg_64) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30511:2 = func.call @_where_476(%30495, %iterArg_65, %iterArg_65) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30512:2 = func.call @_where_477(%30495, %iterArg_66, %iterArg_66) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30513:2 = func.call @_where_478(%30495, %iterArg_67, %iterArg_67) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30514:2 = func.call @_where_479(%30495, %iterArg_68, %iterArg_68) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30515:2 = func.call @_where_480(%30495, %iterArg_69, %iterArg_69) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30516:2 = func.call @_where_481(%30495, %iterArg_70, %iterArg_70) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30517:2 = func.call @_where_482(%30495, %iterArg_71, %iterArg_71) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30518:2 = func.call @_where_483(%30495, %iterArg_72, %iterArg_72) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30519:2 = func.call @_where_484(%30495, %iterArg_73, %iterArg_73) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30520:2 = func.call @_where_485(%30495, %iterArg_74, %iterArg_74) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30521:2 = func.call @_where_486(%30495, %iterArg_75, %iterArg_75) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30522:2 = func.call @_where_487(%30495, %iterArg_76, %iterArg_76) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30523:2 = func.call @_where_488(%30495, %iterArg_77, %iterArg_77) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30524:2 = func.call @_where_489(%30495, %iterArg_78, %iterArg_78) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30525:2 = func.call @_where_490(%30495, %iterArg_79, %iterArg_79) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30526:2 = func.call @_where_491(%30495, %iterArg_80, %iterArg_80) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30527:2 = func.call @_where_492(%30495, %iterArg_81, %iterArg_81) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30528:2 = func.call @_where_493(%30495, %iterArg_82, %iterArg_82) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30529:2 = func.call @_where_494(%30495, %iterArg_83, %iterArg_83) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30530:2 = func.call @_where_495(%30495, %iterArg_84, %iterArg_84) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30531:2 = func.call @_where_496(%30495, %iterArg_85, %iterArg_85) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30532:2 = func.call @_where_497(%30495, %iterArg_86, %iterArg_86) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30533:2 = func.call @_where_498(%30495, %iterArg_87, %iterArg_87) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30534:2 = func.call @_where_499(%30495, %iterArg_88, %iterArg_88) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30535:2 = func.call @_where_500(%30495, %iterArg_89, %iterArg_89) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30536:2 = func.call @_where_501(%30495, %iterArg_90, %iterArg_90) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30537:2 = func.call @_where_502(%30495, %iterArg_91, %iterArg_91) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30538:2 = func.call @_where_503(%30495, %iterArg_92, %iterArg_92) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30539:2 = func.call @_where_504(%30495, %iterArg_93, %iterArg_93) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30540:2 = func.call @_where_505(%30495, %iterArg_94, %iterArg_94) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30541:2 = func.call @_where_506(%30495, %iterArg_95, %iterArg_95) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30542:2 = func.call @_where_507(%30495, %iterArg_96, %iterArg_96) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30543:2 = func.call @_where_508(%30495, %iterArg_97, %iterArg_97) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30544:2 = func.call @_where_509(%30495, %iterArg_98, %iterArg_98) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30545:2 = func.call @_where_510(%30495, %iterArg_99, %iterArg_99) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30546:2 = func.call @_where_511(%30495, %iterArg_100, %iterArg_100) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30547:2 = func.call @_where_512(%30495, %iterArg_101, %iterArg_101) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30548:2 = func.call @_where_513(%30495, %iterArg_102, %iterArg_102) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30549:2 = func.call @_where_514(%30495, %iterArg_103, %iterArg_103) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30550:2 = func.call @_where_515(%30495, %iterArg_104, %iterArg_104) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30551:2 = func.call @_where_516(%30495, %iterArg_105, %iterArg_105) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30552:2 = func.call @_where_517(%30495, %iterArg_106, %iterArg_106) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30553:2 = func.call @_where_518(%30495, %iterArg_107, %iterArg_107) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30554:2 = func.call @_where_519(%30495, %iterArg_108, %iterArg_108) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30555:2 = func.call @_where_520(%30495, %iterArg_109, %iterArg_109) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30556:2 = func.call @_where_521(%30495, %iterArg_110, %iterArg_110) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30557:2 = func.call @_where_522(%30495, %iterArg_111, %iterArg_111) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30558:2 = func.call @_where_523(%30495, %iterArg_112, %iterArg_112) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30559:2 = func.call @_where_524(%30495, %iterArg_113, %iterArg_113) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30560:2 = func.call @_where_525(%30495, %iterArg_114, %iterArg_114) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30561:2 = func.call @_where_526(%30495, %iterArg_115, %iterArg_115) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30562:2 = func.call @_where_527(%30495, %iterArg_116, %iterArg_116) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30563:2 = func.call @_where_528(%30495, %iterArg_117, %iterArg_117) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30564:2 = func.call @_where_529(%30495, %iterArg_118, %iterArg_118) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30565:2 = func.call @_where_530(%30495, %iterArg_119, %iterArg_119) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30566:2 = func.call @_where_531(%30495, %iterArg_120, %iterArg_120) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30567:2 = func.call @_where_532(%30495, %iterArg_121, %iterArg_121) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30568:2 = func.call @_where_533(%30495, %iterArg_122, %iterArg_122) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30569:2 = func.call @_where_534(%30495, %iterArg_123, %iterArg_123) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30570:2 = func.call @_where_535(%30495, %iterArg_124, %iterArg_124) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30571:2 = func.call @_where_536(%30495, %iterArg_125, %iterArg_125) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30572:2 = func.call @_where_537(%30495, %iterArg_126, %iterArg_126) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30573:2 = func.call @_where_538(%30495, %iterArg_127, %iterArg_127) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30574:2 = func.call @_where_539(%30495, %iterArg_128, %iterArg_128) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30575:2 = func.call @_where_540(%30495, %iterArg_129, %iterArg_129) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30576:2 = func.call @_where_541(%30495, %iterArg_130, %iterArg_130) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30577:2 = func.call @_where_542(%30495, %iterArg_131, %iterArg_131) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30578 = func.call @_where_543(%30495) : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc2017)
      %30579:2 = func.call @_where_544(%30495, %iterArg_132, %iterArg_132) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30580:2 = func.call @_where_545(%30495, %iterArg_133, %iterArg_133) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30581:2 = func.call @_where_546(%30495, %iterArg_134, %iterArg_134) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30582:2 = func.call @_where_547(%30495, %iterArg_135, %iterArg_135) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30583:2 = func.call @_where_548(%30495, %iterArg_136, %iterArg_136) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30584:2 = func.call @_where_549(%30495, %iterArg_137, %iterArg_137) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30585:2 = func.call @_where_550(%30495, %iterArg_138, %iterArg_138) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30586:2 = func.call @_where_551(%30495, %iterArg_139, %iterArg_139) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30587:2 = func.call @_where_552(%30495, %iterArg_140, %iterArg_140) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30588:2 = func.call @_where_553(%30495, %iterArg_141, %iterArg_141) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30589:2 = func.call @_where_554(%30495, %iterArg_142, %iterArg_142) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30590:2 = func.call @_where_555(%30495, %iterArg_143, %iterArg_143) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30591:2 = func.call @_where_556(%30495, %iterArg_144, %iterArg_144) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30592:2 = func.call @_where_557(%30495, %iterArg_145, %iterArg_145) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30593:2 = func.call @_where_558(%30495, %iterArg_146, %iterArg_146) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30594:2 = func.call @_where_559(%30495, %iterArg_147, %iterArg_147) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30595:2 = func.call @_where_560(%30495, %iterArg_148, %iterArg_148) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30596:2 = func.call @_where_561(%30495, %iterArg_149, %iterArg_149) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30597:2 = func.call @_where_562(%30495, %iterArg_150, %iterArg_150) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30598:2 = func.call @_where_563(%30495, %iterArg_151, %iterArg_151) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30599:2 = func.call @_where_564(%30495, %iterArg_152, %iterArg_152) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30600:2 = func.call @_where_565(%30495, %iterArg_153, %iterArg_153) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30601:2 = func.call @_where_566(%30495, %iterArg_154, %iterArg_154) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30602:2 = func.call @_where_567(%30495, %iterArg_155, %iterArg_155) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30603:2 = func.call @_where_568(%30495, %iterArg_156, %iterArg_156) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30604:2 = func.call @_where_569(%30495, %iterArg_157, %iterArg_157) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30605:2 = func.call @_where_570(%30495, %iterArg_158, %iterArg_158) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30606:2 = func.call @_where_571(%30495, %iterArg_159, %iterArg_159) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30607:2 = func.call @_where_572(%30495, %iterArg_160, %iterArg_160) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30608:2 = func.call @_where_573(%30495, %iterArg_161, %iterArg_161) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30609:2 = func.call @_where_574(%30495, %iterArg_162, %iterArg_162) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30610:2 = func.call @_where_575(%30495, %iterArg_163, %iterArg_163) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30611:2 = func.call @_where_576(%30495, %iterArg_164, %iterArg_164) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30612:2 = func.call @_where_577(%30495, %iterArg_165, %iterArg_165) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30613:2 = func.call @_where_578(%30495, %iterArg_166, %iterArg_166) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30614:2 = func.call @_where_579(%30495, %iterArg_167, %iterArg_167) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30615:2 = func.call @_where_580(%30495, %iterArg_168, %iterArg_168) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30616:2 = func.call @_where_581(%30495, %iterArg_169, %iterArg_169) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30617:2 = func.call @_where_582(%30495, %iterArg_170, %iterArg_170) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30618:2 = func.call @_where_583(%30495, %iterArg_171, %iterArg_171) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30619:2 = func.call @_where_584(%30495, %iterArg_172, %iterArg_172) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30620:2 = func.call @_where_585(%30495, %iterArg_173, %iterArg_173) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30621:2 = func.call @_where_586(%30495, %iterArg_174, %iterArg_174) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30622:2 = func.call @_where_587(%30495, %iterArg_175, %iterArg_175) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30623:2 = func.call @_where_588(%30495, %iterArg_176, %iterArg_176) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30624:2 = func.call @_where_589(%30495, %iterArg_177, %iterArg_177) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30625:2 = func.call @_where_590(%30495, %iterArg_178, %iterArg_178) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30626:2 = func.call @_where_591(%30495, %iterArg_179, %iterArg_179) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30627:2 = func.call @_where_592(%30495, %iterArg_180, %iterArg_180) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30628:2 = func.call @_where_593(%30495, %iterArg_181, %iterArg_181) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30629:2 = func.call @_where_594(%30495, %iterArg_182, %iterArg_182) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30630:2 = func.call @_where_595(%30495, %iterArg_183, %iterArg_183) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30631:2 = func.call @_where_596(%30495, %iterArg_184, %iterArg_184) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30632:2 = func.call @_where_597(%30495, %iterArg_185, %iterArg_185) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30633:2 = func.call @_where_598(%30495, %iterArg_186, %iterArg_186) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30634:2 = func.call @_where_599(%30495, %iterArg_187, %iterArg_187) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30635:2 = func.call @_where_600(%30495, %iterArg_188, %iterArg_188) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30636:2 = func.call @_where_601(%30495, %iterArg_189, %iterArg_189) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30637:2 = func.call @_where_602(%30495, %iterArg_190, %iterArg_190) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30638:2 = func.call @_where_603(%30495, %iterArg_191, %iterArg_191) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30639:2 = func.call @_where_604(%30495, %iterArg_192, %iterArg_192) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30640:2 = func.call @_where_605(%30495, %iterArg_193, %iterArg_193) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30641:2 = func.call @_where_606(%30495, %iterArg_194, %iterArg_194) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30642:2 = func.call @_where_607(%30495, %iterArg_195, %iterArg_195) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30643:2 = func.call @_where_608(%30495, %iterArg_196, %iterArg_196) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30644:2 = func.call @_where_609(%30495, %iterArg_197, %iterArg_197) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30645:2 = func.call @_where_610(%30495, %iterArg_198, %iterArg_198) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30646:2 = func.call @_where_611(%30495, %iterArg_199, %iterArg_199) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30647:2 = func.call @_where_612(%30495, %iterArg_200, %iterArg_200) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30648:2 = func.call @_where_613(%30495, %iterArg_201, %iterArg_201) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30649:2 = func.call @_where_614(%30495, %iterArg_202, %iterArg_202) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30650:2 = func.call @_where_615(%30495, %iterArg_203, %iterArg_203) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30651:2 = func.call @_where_616(%30495, %iterArg_204, %iterArg_204) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30652:2 = func.call @_where_617(%30495, %iterArg_205, %iterArg_205) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30653:2 = func.call @_where_618(%30495, %iterArg_206, %iterArg_206) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30654:2 = func.call @_where_619(%30495, %iterArg_207, %iterArg_207) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30655:2 = func.call @_where_620(%30495, %iterArg_208, %iterArg_208) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30656:2 = func.call @_where_621(%30495, %iterArg_209, %iterArg_209) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30657:2 = func.call @_where_622(%30495, %iterArg_210, %iterArg_210) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30658:2 = func.call @_where_623(%30495, %iterArg_211, %iterArg_211) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30659:2 = func.call @_where_624(%30495, %iterArg_212, %iterArg_212) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30660:2 = func.call @_where_625(%30495, %iterArg_213, %iterArg_213) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30661:2 = func.call @_where_626(%30495, %iterArg_214, %iterArg_214) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30662:2 = func.call @_where_627(%30495, %iterArg_215, %iterArg_215) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30663:2 = func.call @_where_628(%30495, %iterArg_216, %iterArg_216) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30664:2 = func.call @_where_629(%30495, %iterArg_217, %iterArg_217) : (tensor<1xi1>, tensor<1x8192xbf16>, tensor<1x8192xbf16>) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) loc(#loc2016)
      %30665:2 = func.call @_where_630(%30495, %iterArg_218, %iterArg_218) : (tensor<1xi1>, tensor<1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) loc(#loc2016)
      %30666:2 = func.call @_where_631(%30495, %iterArg_219, %iterArg_219) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30667:2 = func.call @_where_632(%30495, %iterArg_220, %iterArg_220) : (tensor<1xi1>, tensor<1x8192x2048xbf16>, tensor<1x8192x2048xbf16>) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) loc(#loc2016)
      %30668:2 = func.call @_where_633(%30495, %iterArg_221, %iterArg_221) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30669:2 = func.call @_where_634(%30495, %iterArg_222, %iterArg_222) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30670:2 = func.call @_where_635(%30495, %iterArg_223, %iterArg_223) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30671:2 = func.call @_where_636(%30495, %iterArg_224, %iterArg_224) : (tensor<1xi1>, tensor<1x2048xbf16>, tensor<1x2048xbf16>) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) loc(#loc2016)
      %30672:2 = func.call @_where_637(%30495, %iterArg_225, %iterArg_225) : (tensor<1xi1>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xbf16>) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) loc(#loc2016)
      %30673:2 = func.call @_where_638(%30495, %iterArg_226, %iterArg_226) : (tensor<1xi1>, tensor<1x64xbf16>, tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xi1>) loc(#loc2016)
      %30674:2 = func.call @_where_639(%30495, %iterArg_227, %iterArg_227) : (tensor<1xi1>, tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) loc(#loc2016)
      %30675 = stablehlo.iota dim = 0 : tensor<2048xi32> loc(#loc2018)
      %30676 = stablehlo.broadcast_in_dim %30675, dims = [1] : (tensor<2048xi32>) -> tensor<1x2048xi32> loc(#loc2019)
      %30677 = stablehlo.reshape %30676 : (tensor<1x2048xi32>) -> tensor<1x1x1x2048xi32> loc(#loc2020)
      %30678 = stablehlo.reshape %30677 : (tensor<1x1x1x2048xi32>) -> tensor<1x1x2048xi32> loc(#loc2021)
      %30679 = stablehlo.broadcast_in_dim %30678, dims = [1, 2, 3] : (tensor<1x1x2048xi32>) -> tensor<2048x1x1x2048xi32> loc(#loc2022)
      %30680 = stablehlo.reshape %30679 : (tensor<2048x1x1x2048xi32>) -> tensor<2048x2048xi32> loc(#loc2023)
      %30681 = stablehlo.iota dim = 0 : tensor<2048xi32> loc(#loc2024)
      %30682 = stablehlo.broadcast_in_dim %30681, dims = [0] : (tensor<2048xi32>) -> tensor<2048x1xi32> loc(#loc2025)
      %30683 = stablehlo.reshape %30682 : (tensor<2048x1xi32>) -> tensor<1x2048x1x1xi32> loc(#loc2026)
      %30684 = stablehlo.reshape %30683 : (tensor<1x2048x1x1xi32>) -> tensor<1x2048x1xi32> loc(#loc2027)
      %30685 = stablehlo.broadcast_in_dim %30684, dims = [0, 1, 3] : (tensor<1x2048x1xi32>) -> tensor<1x2048x2048x1xi32> loc(#loc2028)
      %30686 = stablehlo.reshape %30685 : (tensor<1x2048x2048x1xi32>) -> tensor<2048x2048xi32> loc(#loc2029)
      %30687 = stablehlo.compare  LT, %30686, %30680,  SIGNED : (tensor<2048x2048xi32>, tensor<2048x2048xi32>) -> tensor<2048x2048xi1> loc(#loc2030)
      %30688 = stablehlo.convert %30687 : (tensor<2048x2048xi1>) -> tensor<2048x2048xbf16> loc(#loc2031)
      %30689 = stablehlo.constant dense<-2.366030e+38> : tensor<bf16> loc(#loc1984)
      %30690 = stablehlo.broadcast_in_dim %30689, dims = [] : (tensor<bf16>) -> tensor<2048x2048xbf16> loc(#loc2032)
      %30691 = stablehlo.multiply %30688, %30690 : tensor<2048x2048xbf16> loc(#loc2032)
      %30692 = stablehlo.broadcast_in_dim %30691, dims = [2, 3] : (tensor<2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2033)
      %30693 = stablehlo.broadcast_in_dim %30692, dims = [1, 2, 3, 4] : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc2034)
      %30694 = stablehlo.minimum %30498, %30693 : tensor<1x1x1x2048x2048xbf16> loc(#loc2035)
      %30695 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2036)
      %30696 = stablehlo.convert %30496 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2037)
      %30697 = stablehlo.convert %30695 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2038)
      %30698 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30699 = stablehlo.broadcast_in_dim %30698, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2039)
      %30700 = stablehlo.subtract %30699, %30697 : tensor<1x1x2048x1xf32> loc(#loc2039)
      %30701 = stablehlo.convert %30700 : tensor<1x1x2048x1xf32> loc(#loc2040)
      %30702 = stablehlo.broadcast_in_dim %30701, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30703 = stablehlo.multiply %30696, %30702 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30704 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2042)
      %30705 = stablehlo.reduce(%30703 init: %30704) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2042)
        stablehlo.return %39096 : tensor<f32> loc(#loc2042)
      } loc(#loc2042)
      %30706 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30707 = stablehlo.broadcast_in_dim %30706, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2043)
      %30708 = stablehlo.convert %30700 : tensor<1x1x2048x1xf32> loc(#loc2040)
      %30709 = stablehlo.broadcast_in_dim %30707, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2044)
      %30710 = stablehlo.broadcast_in_dim %30708, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30711 = stablehlo.multiply %30709, %30710 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30712 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2042)
      %30713 = stablehlo.reduce(%30711 init: %30712) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2042)
        stablehlo.return %39096 : tensor<f32> loc(#loc2042)
      } loc(#loc2042)
      %30714 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30715 = stablehlo.broadcast_in_dim %30714, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2045)
      %30716 = stablehlo.maximum %30715, %30713 : tensor<1xf32> loc(#loc2045)
      %30717 = stablehlo.divide %30705, %30716 : tensor<1xf32> loc(#loc2046)
      %30718 = stablehlo.broadcast_in_dim %30717, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2047)
      %30719 = stablehlo.broadcast_in_dim %30718, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2048)
      %30720 = stablehlo.subtract %30696, %30719 : tensor<1x1x2048x2048xf32> loc(#loc2048)
      %30721 = stablehlo.convert %30700 : tensor<1x1x2048x1xf32> loc(#loc2040)
      %30722 = stablehlo.broadcast_in_dim %30721, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30723 = stablehlo.multiply %30720, %30722 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30724 = stablehlo.multiply %30723, %30723 : tensor<1x1x2048x2048xf32> loc(#loc2049)
      %30725 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %30726 = stablehlo.broadcast_in_dim %30725, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2049)
      %30727 = stablehlo.multiply %30726, %30723 : tensor<1x1x2048x2048xf32> loc(#loc2049)
      %30728 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2042)
      %30729 = stablehlo.reduce(%30724 init: %30728) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2042)
        stablehlo.return %39096 : tensor<f32> loc(#loc2042)
      } loc(#loc2042)
      %30730 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30731 = stablehlo.broadcast_in_dim %30730, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2045)
      %30732 = stablehlo.maximum %30731, %30713 : tensor<1xf32> loc(#loc2045)
      %30733 = stablehlo.divide %30729, %30732 : tensor<1xf32> loc(#loc2046)
      %30734 = stablehlo.sqrt %30733 : tensor<1xf32> loc(#loc2050)
      %30735 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %30736 = stablehlo.broadcast_in_dim %30735, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2051)
      %30737 = stablehlo.divide %30736, %30734 : tensor<1xf32> loc(#loc2051)
      %30738 = stablehlo.convert %30700 : tensor<1x1x2048x1xf32> loc(#loc2040)
      %30739 = stablehlo.broadcast_in_dim %30738, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30740 = stablehlo.multiply %30696, %30739 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %30741 = stablehlo.abs %30740 : tensor<1x1x2048x2048xf32> loc(#loc2052)
      %30742 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %30743 = stablehlo.broadcast_in_dim %30742, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2053)
      %30744 = stablehlo.compare  GE, %30740, %30743,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2053)
      %30745 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2054)
      %30746 = stablehlo.reduce(%30741 init: %30745) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2054)
        stablehlo.return %39096 : tensor<f32> loc(#loc2054)
      } loc(#loc2054)
      %30747 = stablehlo.reshape %30746 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2055)
      %30748 = stablehlo.broadcast_in_dim %30747, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2056)
      %30749 = stablehlo.compare  EQ, %30741, %30748,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2056)
      %30750 = stablehlo.convert %30749 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2057)
      %30751 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2058)
      %30752 = stablehlo.reduce(%30750 init: %30751) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2058)
        stablehlo.return %39096 : tensor<f32> loc(#loc2058)
      } loc(#loc2058)
      %30753 = stablehlo.convert %30496 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2059)
      %30754 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2060)
      %30755 = stablehlo.reduce(%30753 init: %30754) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2060)
        stablehlo.return %39096 : tensor<f32> loc(#loc2060)
      } loc(#loc2060)
      %30756 = stablehlo.broadcast_in_dim %30755, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2061)
      %30757 = stablehlo.convert %30756 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2062)
      %30758 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %30759 = stablehlo.broadcast_in_dim %30758, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2063)
      %30760 = stablehlo.divide %30757, %30759 : tensor<1x1x2048x1xbf16> loc(#loc2063)
      %30761 = stablehlo.broadcast_in_dim %30760, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2064)
      %30762 = stablehlo.subtract %30496, %30761 : tensor<1x1x2048x2048xbf16> loc(#loc2064)
      %30763 = stablehlo.multiply %30762, %30762 : tensor<1x1x2048x2048xbf16> loc(#loc2065)
      %30764 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %30765 = stablehlo.broadcast_in_dim %30764, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2065)
      %30766 = stablehlo.multiply %30765, %30762 : tensor<1x1x2048x2048xbf16> loc(#loc2065)
      %30767 = stablehlo.convert %30763 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2059)
      %30768 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2060)
      %30769 = stablehlo.reduce(%30767 init: %30768) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2060)
        stablehlo.return %39096 : tensor<f32> loc(#loc2060)
      } loc(#loc2060)
      %30770 = stablehlo.broadcast_in_dim %30769, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2061)
      %30771 = stablehlo.convert %30770 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2062)
      %30772 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %30773 = stablehlo.broadcast_in_dim %30772, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2063)
      %30774 = stablehlo.divide %30771, %30773 : tensor<1x1x2048x1xbf16> loc(#loc2063)
      %30775 = stablehlo.broadcast_in_dim %30760, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2064)
      %30776 = stablehlo.subtract %30496, %30775 : tensor<1x1x2048x2048xbf16> loc(#loc2064)
      %30777 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %30778 = stablehlo.broadcast_in_dim %30777, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2066)
      %30779 = stablehlo.add %30774, %30778 : tensor<1x1x2048x1xbf16> loc(#loc2066)
      %30780 = stablehlo.rsqrt %30779 : tensor<1x1x2048x1xbf16> loc(#loc2067)
      %30781 = stablehlo.divide %30780, %30779 : tensor<1x1x2048x1xbf16> loc(#loc2068)
      %30782 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %30783 = stablehlo.broadcast_in_dim %30782, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2069)
      %30784 = stablehlo.multiply %30783, %30781 : tensor<1x1x2048x1xbf16> loc(#loc2069)
      %30785 = stablehlo.broadcast_in_dim %30780, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2069)
      %30786 = stablehlo.multiply %30776, %30785 : tensor<1x1x2048x2048xbf16> loc(#loc2069)
      %30787 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %30788 = stablehlo.broadcast_in_dim %30787, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2070)
      %30789 = stablehlo.add %30788, %30506#0 : tensor<1x2048xbf16> loc(#loc2070)
      %30790 = stablehlo.broadcast_in_dim %30789, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2071)
      %30791 = stablehlo.broadcast_in_dim %30790, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2072)
      %30792 = stablehlo.multiply %30786, %30791 : tensor<1x1x2048x2048xbf16> loc(#loc2072)
      %30793 = stablehlo.broadcast_in_dim %30505#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2073)
      %30794 = stablehlo.broadcast_in_dim %30793, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2074)
      %30795 = stablehlo.add %30792, %30794 : tensor<1x1x2048x2048xbf16> loc(#loc2074)
      %30796 = func.call @_einsum_640(%30795, %30507#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2075)
      %30797 = "stablehlo.slice"(%30796) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2076)
      %30798 = stablehlo.reshape %30797 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2077)
      %30799 = "stablehlo.slice"(%30796) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2078)
      %30800 = stablehlo.reshape %30799 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2077)
      %30801 = "stablehlo.slice"(%30796) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2079)
      %30802 = stablehlo.reshape %30801 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2077)
      %30803 = arith.select %ptrue, %30798, %30798 : tensor<1x1x2048x32x64xbf16>
      %30804 = arith.select %ptrue, %30800, %30800 : tensor<1x1x2048x32x64xbf16>
      %30805 = arith.select %ptrue, %30802, %30802 : tensor<1x1x2048x32x64xbf16>
      %30806:2 = func.call @softplus_641(%30508#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2081)
      %30807 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %30808 = stablehlo.broadcast_in_dim %30807, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2082)
      %30809 = stablehlo.multiply %30808, %30806#0 : tensor<1x64xbf16> loc(#loc2082)
      %30810 = stablehlo.broadcast_in_dim %30809, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2083)
      %30811 = stablehlo.broadcast_in_dim %30810, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2084)
      %30812 = stablehlo.multiply %30803, %30811 : tensor<1x1x2048x32x64xbf16> loc(#loc2084)
      %30813 = func.call @_einsum_643(%30812, %30804) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2085)
      %30814 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %30815 = stablehlo.broadcast_in_dim %30814, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2086)
      %30816 = stablehlo.divide %30813, %30815 : tensor<1x1x32x2048x2048xbf16> loc(#loc2086)
      %30817 = stablehlo.tanh %30816 : tensor<1x1x32x2048x2048xbf16> loc(#loc2087)
      %30818 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %30819 = stablehlo.broadcast_in_dim %30818, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2088)
      %30820 = stablehlo.subtract %30819, %30817 : tensor<1x1x32x2048x2048xbf16> loc(#loc2088)
      %30821 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %30822 = stablehlo.broadcast_in_dim %30821, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2089)
      %30823 = stablehlo.multiply %30822, %30817 : tensor<1x1x32x2048x2048xbf16> loc(#loc2089)
      %30824 = stablehlo.convert %30823 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2090)
      %30825 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %30826 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %30827 = stablehlo.multiply %30825, %30826 : tensor<f32> loc(#loc2091)
      %30828 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2092)
      %30829 = stablehlo.broadcast_in_dim %30827, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2093)
      %30830 = stablehlo.compare  GE, %30828, %30829,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2093)
      %30831 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %30832:2 = func.call @_where_644(%30830, %30824, %30831) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2094)
      %30833 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2095)
      %30834 = stablehlo.reduce(%30832#0 init: %30833) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2095)
        stablehlo.return %39096 : tensor<f32> loc(#loc2095)
      } loc(#loc2095)
      %30835 = stablehlo.broadcast_in_dim %30834, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2096)
      %30836 = stablehlo.broadcast_in_dim %30835, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2097)
      %30837 = stablehlo.subtract %30832#0, %30836 : tensor<1x1x32x2048x2048xf32> loc(#loc2097)
      %30838 = stablehlo.exponential %30837 : tensor<1x1x32x2048x2048xf32> loc(#loc2098)
      %30839 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2099)
      %30840 = stablehlo.reduce(%30838 init: %30839) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2099)
        stablehlo.return %39096 : tensor<f32> loc(#loc2099)
      } loc(#loc2099)
      %30841 = stablehlo.broadcast_in_dim %30840, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2096)
      %30842 = stablehlo.broadcast_in_dim %30841, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2100)
      %30843 = stablehlo.divide %30838, %30842 : tensor<1x1x32x2048x2048xf32> loc(#loc2100)
      %30844 = stablehlo.multiply %30841, %30841 : tensor<1x1x32x2048x1xf32> loc(#loc2101)
      %30845 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2102)
      %30846 = stablehlo.broadcast_in_dim %30845, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2103)
      %30847 = stablehlo.divide %30846, %30844 : tensor<1x1x32x2048x1xf32> loc(#loc2100)
      %30848 = stablehlo.convert %30843 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2104)
      %30849 = func.call @_einsum_645(%30848, %30805) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2105)
      %30850 = arith.select %ptrue, %30849, %30849 : tensor<1x1x2048x32x64xbf16>
      %30851 = func.call @_einsum_646(%30850, %30509#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2106)
      %30852 = arith.select %ptrue, %30851, %30851 : tensor<1x1x2048x2048xbf16>
      %30853 = stablehlo.add %30852, %30496 : tensor<1x1x2048x2048xbf16> loc(#loc2108)
      %30854 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2109)
      %30855 = stablehlo.convert %30853 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2110)
      %30856 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2111)
      %30857 = stablehlo.reduce(%30855 init: %30856) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2111)
        stablehlo.return %39096 : tensor<f32> loc(#loc2111)
      } loc(#loc2111)
      %30858 = stablehlo.broadcast_in_dim %30857, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2112)
      %30859 = stablehlo.convert %30858 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2113)
      %30860 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %30861 = stablehlo.broadcast_in_dim %30860, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2114)
      %30862 = stablehlo.divide %30859, %30861 : tensor<1x1x2048x1xbf16> loc(#loc2114)
      %30863 = stablehlo.broadcast_in_dim %30862, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2115)
      %30864 = stablehlo.subtract %30853, %30863 : tensor<1x1x2048x2048xbf16> loc(#loc2115)
      %30865 = stablehlo.multiply %30864, %30864 : tensor<1x1x2048x2048xbf16> loc(#loc2116)
      %30866 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %30867 = stablehlo.broadcast_in_dim %30866, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2116)
      %30868 = stablehlo.multiply %30867, %30864 : tensor<1x1x2048x2048xbf16> loc(#loc2116)
      %30869 = stablehlo.convert %30865 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2110)
      %30870 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2111)
      %30871 = stablehlo.reduce(%30869 init: %30870) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2111)
        stablehlo.return %39096 : tensor<f32> loc(#loc2111)
      } loc(#loc2111)
      %30872 = stablehlo.broadcast_in_dim %30871, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2112)
      %30873 = stablehlo.convert %30872 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2113)
      %30874 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %30875 = stablehlo.broadcast_in_dim %30874, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2114)
      %30876 = stablehlo.divide %30873, %30875 : tensor<1x1x2048x1xbf16> loc(#loc2114)
      %30877 = stablehlo.broadcast_in_dim %30862, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2115)
      %30878 = stablehlo.subtract %30853, %30877 : tensor<1x1x2048x2048xbf16> loc(#loc2115)
      %30879 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %30880 = stablehlo.broadcast_in_dim %30879, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2117)
      %30881 = stablehlo.add %30876, %30880 : tensor<1x1x2048x1xbf16> loc(#loc2117)
      %30882 = stablehlo.rsqrt %30881 : tensor<1x1x2048x1xbf16> loc(#loc2118)
      %30883 = stablehlo.divide %30882, %30881 : tensor<1x1x2048x1xbf16> loc(#loc2119)
      %30884 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %30885 = stablehlo.broadcast_in_dim %30884, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2120)
      %30886 = stablehlo.multiply %30885, %30883 : tensor<1x1x2048x1xbf16> loc(#loc2120)
      %30887 = stablehlo.broadcast_in_dim %30882, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2120)
      %30888 = stablehlo.multiply %30878, %30887 : tensor<1x1x2048x2048xbf16> loc(#loc2120)
      %30889 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %30890 = stablehlo.broadcast_in_dim %30889, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2121)
      %30891 = stablehlo.add %30890, %30504#0 : tensor<1x2048xbf16> loc(#loc2121)
      %30892 = stablehlo.broadcast_in_dim %30891, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2122)
      %30893 = stablehlo.broadcast_in_dim %30892, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2123)
      %30894 = stablehlo.multiply %30888, %30893 : tensor<1x1x2048x2048xbf16> loc(#loc2123)
      %30895 = stablehlo.broadcast_in_dim %30503#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2124)
      %30896 = stablehlo.broadcast_in_dim %30895, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2125)
      %30897 = stablehlo.add %30894, %30896 : tensor<1x1x2048x2048xbf16> loc(#loc2125)
      %30898 = func.call @_einsum_647(%30897, %30500#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2126)
      %30899 = arith.select %ptrue, %30898, %30898 : tensor<1x1x2048x8192xbf16>
      %30900 = stablehlo.broadcast_in_dim %30499#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2128)
      %30901 = stablehlo.broadcast_in_dim %30900, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2129)
      %30902 = stablehlo.add %30899, %30901 : tensor<1x1x2048x8192xbf16> loc(#loc2129)
      %30903 = stablehlo.multiply %30902, %30902 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30904 = stablehlo.multiply %30902, %30903 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30905 = stablehlo.multiply %30902, %30902 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30906 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %30907 = stablehlo.broadcast_in_dim %30906, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30908 = stablehlo.multiply %30907, %30905 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30909 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %30910 = stablehlo.broadcast_in_dim %30909, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30911 = stablehlo.multiply %30910, %30904 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30912 = stablehlo.add %30902, %30911 : tensor<1x1x2048x8192xbf16> loc(#loc2131)
      %30913 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %30914 = stablehlo.broadcast_in_dim %30913, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30915 = stablehlo.multiply %30914, %30912 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30916 = stablehlo.tanh %30915 : tensor<1x1x2048x8192xbf16> loc(#loc2132)
      %30917 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %30918 = stablehlo.broadcast_in_dim %30917, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2133)
      %30919 = stablehlo.subtract %30918, %30916 : tensor<1x1x2048x8192xbf16> loc(#loc2133)
      %30920 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %30921 = stablehlo.broadcast_in_dim %30920, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2131)
      %30922 = stablehlo.add %30921, %30916 : tensor<1x1x2048x8192xbf16> loc(#loc2131)
      %30923 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %30924 = stablehlo.broadcast_in_dim %30923, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30925 = stablehlo.multiply %30924, %30922 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30926 = stablehlo.multiply %30902, %30925 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %30927 = stablehlo.convert %30854 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2134)
      %30928 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30929 = stablehlo.broadcast_in_dim %30928, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2135)
      %30930 = stablehlo.subtract %30929, %30927 : tensor<1x1x2048x1xf32> loc(#loc2135)
      %30931 = stablehlo.convert %30930 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2136)
      %30932 = stablehlo.broadcast_in_dim %30931, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2137)
      %30933 = stablehlo.multiply %30926, %30932 : tensor<1x1x2048x8192xbf16> loc(#loc2137)
      %30934 = func.call @_einsum_648(%30933, %30502#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2138)
      %30935 = arith.select %ptrue, %30934, %30934 : tensor<1x1x2048x2048xbf16>
      %30936 = stablehlo.broadcast_in_dim %30501#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2140)
      %30937 = stablehlo.broadcast_in_dim %30936, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2141)
      %30938 = stablehlo.add %30935, %30937 : tensor<1x1x2048x2048xbf16> loc(#loc2141)
      %30939 = stablehlo.convert %30854 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2134)
      %30940 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30941 = stablehlo.broadcast_in_dim %30940, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2135)
      %30942 = stablehlo.subtract %30941, %30939 : tensor<1x1x2048x1xf32> loc(#loc2135)
      %30943 = stablehlo.convert %30942 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2142)
      %30944 = stablehlo.broadcast_in_dim %30943, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2143)
      %30945 = stablehlo.multiply %30938, %30944 : tensor<1x1x2048x2048xbf16> loc(#loc2143)
      %30946 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %30947 = stablehlo.broadcast_in_dim %30946, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2144)
      %30948 = stablehlo.multiply %30945, %30947 : tensor<1x1x2048x2048xbf16> loc(#loc2144)
      %30949 = stablehlo.add %30853, %30948 : tensor<1x1x2048x2048xbf16> loc(#loc2145)
      %30950 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2146)
      %30951 = stablehlo.convert %30949 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2147)
      %30952 = stablehlo.convert %30950 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2148)
      %30953 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30954 = stablehlo.broadcast_in_dim %30953, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2149)
      %30955 = stablehlo.subtract %30954, %30952 : tensor<1x1x2048x1xf32> loc(#loc2149)
      %30956 = stablehlo.convert %30955 : tensor<1x1x2048x1xf32> loc(#loc2150)
      %30957 = stablehlo.broadcast_in_dim %30956, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30958 = stablehlo.multiply %30951, %30957 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30959 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2152)
      %30960 = stablehlo.reduce(%30958 init: %30959) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2152)
        stablehlo.return %39096 : tensor<f32> loc(#loc2152)
      } loc(#loc2152)
      %30961 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30962 = stablehlo.broadcast_in_dim %30961, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2153)
      %30963 = stablehlo.convert %30955 : tensor<1x1x2048x1xf32> loc(#loc2150)
      %30964 = stablehlo.broadcast_in_dim %30962, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2154)
      %30965 = stablehlo.broadcast_in_dim %30963, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30966 = stablehlo.multiply %30964, %30965 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30967 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2152)
      %30968 = stablehlo.reduce(%30966 init: %30967) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2152)
        stablehlo.return %39096 : tensor<f32> loc(#loc2152)
      } loc(#loc2152)
      %30969 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30970 = stablehlo.broadcast_in_dim %30969, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2155)
      %30971 = stablehlo.maximum %30970, %30968 : tensor<1xf32> loc(#loc2155)
      %30972 = stablehlo.divide %30960, %30971 : tensor<1xf32> loc(#loc2156)
      %30973 = stablehlo.broadcast_in_dim %30972, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2157)
      %30974 = stablehlo.broadcast_in_dim %30973, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2158)
      %30975 = stablehlo.subtract %30951, %30974 : tensor<1x1x2048x2048xf32> loc(#loc2158)
      %30976 = stablehlo.convert %30955 : tensor<1x1x2048x1xf32> loc(#loc2150)
      %30977 = stablehlo.broadcast_in_dim %30976, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30978 = stablehlo.multiply %30975, %30977 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30979 = stablehlo.multiply %30978, %30978 : tensor<1x1x2048x2048xf32> loc(#loc2159)
      %30980 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %30981 = stablehlo.broadcast_in_dim %30980, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2159)
      %30982 = stablehlo.multiply %30981, %30978 : tensor<1x1x2048x2048xf32> loc(#loc2159)
      %30983 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2152)
      %30984 = stablehlo.reduce(%30979 init: %30983) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2152)
        stablehlo.return %39096 : tensor<f32> loc(#loc2152)
      } loc(#loc2152)
      %30985 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %30986 = stablehlo.broadcast_in_dim %30985, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2155)
      %30987 = stablehlo.maximum %30986, %30968 : tensor<1xf32> loc(#loc2155)
      %30988 = stablehlo.divide %30984, %30987 : tensor<1xf32> loc(#loc2156)
      %30989 = stablehlo.sqrt %30988 : tensor<1xf32> loc(#loc2160)
      %30990 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %30991 = stablehlo.broadcast_in_dim %30990, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2161)
      %30992 = stablehlo.divide %30991, %30989 : tensor<1xf32> loc(#loc2161)
      %30993 = stablehlo.convert %30955 : tensor<1x1x2048x1xf32> loc(#loc2150)
      %30994 = stablehlo.broadcast_in_dim %30993, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30995 = stablehlo.multiply %30951, %30994 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %30996 = stablehlo.abs %30995 : tensor<1x1x2048x2048xf32> loc(#loc2162)
      %30997 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %30998 = stablehlo.broadcast_in_dim %30997, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2163)
      %30999 = stablehlo.compare  GE, %30995, %30998,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2163)
      %31000 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2164)
      %31001 = stablehlo.reduce(%30996 init: %31000) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2164)
        stablehlo.return %39096 : tensor<f32> loc(#loc2164)
      } loc(#loc2164)
      %31002 = stablehlo.reshape %31001 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2165)
      %31003 = stablehlo.broadcast_in_dim %31002, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2166)
      %31004 = stablehlo.compare  EQ, %30996, %31003,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2166)
      %31005 = stablehlo.convert %31004 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2167)
      %31006 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2168)
      %31007 = stablehlo.reduce(%31005 init: %31006) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2168)
        stablehlo.return %39096 : tensor<f32> loc(#loc2168)
      } loc(#loc2168)
      %31008 = stablehlo.convert %30949 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2169)
      %31009 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2170)
      %31010 = stablehlo.reduce(%31008 init: %31009) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2170)
        stablehlo.return %39096 : tensor<f32> loc(#loc2170)
      } loc(#loc2170)
      %31011 = stablehlo.broadcast_in_dim %31010, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2171)
      %31012 = stablehlo.convert %31011 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2172)
      %31013 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31014 = stablehlo.broadcast_in_dim %31013, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2173)
      %31015 = stablehlo.divide %31012, %31014 : tensor<1x1x2048x1xbf16> loc(#loc2173)
      %31016 = stablehlo.broadcast_in_dim %31015, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2174)
      %31017 = stablehlo.subtract %30949, %31016 : tensor<1x1x2048x2048xbf16> loc(#loc2174)
      %31018 = stablehlo.multiply %31017, %31017 : tensor<1x1x2048x2048xbf16> loc(#loc2175)
      %31019 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31020 = stablehlo.broadcast_in_dim %31019, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2175)
      %31021 = stablehlo.multiply %31020, %31017 : tensor<1x1x2048x2048xbf16> loc(#loc2175)
      %31022 = stablehlo.convert %31018 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2169)
      %31023 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2170)
      %31024 = stablehlo.reduce(%31022 init: %31023) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2170)
        stablehlo.return %39096 : tensor<f32> loc(#loc2170)
      } loc(#loc2170)
      %31025 = stablehlo.broadcast_in_dim %31024, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2171)
      %31026 = stablehlo.convert %31025 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2172)
      %31027 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31028 = stablehlo.broadcast_in_dim %31027, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2173)
      %31029 = stablehlo.divide %31026, %31028 : tensor<1x1x2048x1xbf16> loc(#loc2173)
      %31030 = stablehlo.broadcast_in_dim %31015, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2174)
      %31031 = stablehlo.subtract %30949, %31030 : tensor<1x1x2048x2048xbf16> loc(#loc2174)
      %31032 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31033 = stablehlo.broadcast_in_dim %31032, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2176)
      %31034 = stablehlo.add %31029, %31033 : tensor<1x1x2048x1xbf16> loc(#loc2176)
      %31035 = stablehlo.rsqrt %31034 : tensor<1x1x2048x1xbf16> loc(#loc2177)
      %31036 = stablehlo.divide %31035, %31034 : tensor<1x1x2048x1xbf16> loc(#loc2178)
      %31037 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31038 = stablehlo.broadcast_in_dim %31037, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2179)
      %31039 = stablehlo.multiply %31038, %31036 : tensor<1x1x2048x1xbf16> loc(#loc2179)
      %31040 = stablehlo.broadcast_in_dim %31035, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2179)
      %31041 = stablehlo.multiply %31031, %31040 : tensor<1x1x2048x2048xbf16> loc(#loc2179)
      %31042 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31043 = stablehlo.broadcast_in_dim %31042, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2180)
      %31044 = stablehlo.add %31043, %30517#0 : tensor<1x2048xbf16> loc(#loc2180)
      %31045 = stablehlo.broadcast_in_dim %31044, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2181)
      %31046 = stablehlo.broadcast_in_dim %31045, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2182)
      %31047 = stablehlo.multiply %31041, %31046 : tensor<1x1x2048x2048xbf16> loc(#loc2182)
      %31048 = stablehlo.broadcast_in_dim %30516#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2183)
      %31049 = stablehlo.broadcast_in_dim %31048, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2184)
      %31050 = stablehlo.add %31047, %31049 : tensor<1x1x2048x2048xbf16> loc(#loc2184)
      %31051 = func.call @_einsum_649(%31050, %30518#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2185)
      %31052 = "stablehlo.slice"(%31051) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2186)
      %31053 = stablehlo.reshape %31052 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2187)
      %31054 = "stablehlo.slice"(%31051) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2188)
      %31055 = stablehlo.reshape %31054 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2187)
      %31056 = "stablehlo.slice"(%31051) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2189)
      %31057 = stablehlo.reshape %31056 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2187)
      %31058 = arith.select %ptrue, %31053, %31053 : tensor<1x1x2048x32x64xbf16>
      %31059 = arith.select %ptrue, %31055, %31055 : tensor<1x1x2048x32x64xbf16>
      %31060 = arith.select %ptrue, %31057, %31057 : tensor<1x1x2048x32x64xbf16>
      %31061:2 = func.call @softplus_650(%30519#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2191)
      %31062 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %31063 = stablehlo.broadcast_in_dim %31062, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2192)
      %31064 = stablehlo.multiply %31063, %31061#0 : tensor<1x64xbf16> loc(#loc2192)
      %31065 = stablehlo.broadcast_in_dim %31064, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2193)
      %31066 = stablehlo.broadcast_in_dim %31065, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2194)
      %31067 = stablehlo.multiply %31058, %31066 : tensor<1x1x2048x32x64xbf16> loc(#loc2194)
      %31068 = func.call @_einsum_652(%31067, %31059) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2195)
      %31069 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31070 = stablehlo.broadcast_in_dim %31069, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2196)
      %31071 = stablehlo.divide %31068, %31070 : tensor<1x1x32x2048x2048xbf16> loc(#loc2196)
      %31072 = stablehlo.tanh %31071 : tensor<1x1x32x2048x2048xbf16> loc(#loc2197)
      %31073 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31074 = stablehlo.broadcast_in_dim %31073, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2198)
      %31075 = stablehlo.subtract %31074, %31072 : tensor<1x1x32x2048x2048xbf16> loc(#loc2198)
      %31076 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31077 = stablehlo.broadcast_in_dim %31076, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2199)
      %31078 = stablehlo.multiply %31077, %31072 : tensor<1x1x32x2048x2048xbf16> loc(#loc2199)
      %31079 = stablehlo.convert %31078 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2200)
      %31080 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31081 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31082 = stablehlo.multiply %31080, %31081 : tensor<f32> loc(#loc2201)
      %31083 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2202)
      %31084 = stablehlo.broadcast_in_dim %31082, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2203)
      %31085 = stablehlo.compare  GE, %31083, %31084,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2203)
      %31086 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31087:2 = func.call @_where_653(%31085, %31079, %31086) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2204)
      %31088 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2205)
      %31089 = stablehlo.reduce(%31087#0 init: %31088) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2205)
        stablehlo.return %39096 : tensor<f32> loc(#loc2205)
      } loc(#loc2205)
      %31090 = stablehlo.broadcast_in_dim %31089, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2206)
      %31091 = stablehlo.broadcast_in_dim %31090, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2207)
      %31092 = stablehlo.subtract %31087#0, %31091 : tensor<1x1x32x2048x2048xf32> loc(#loc2207)
      %31093 = stablehlo.exponential %31092 : tensor<1x1x32x2048x2048xf32> loc(#loc2208)
      %31094 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2209)
      %31095 = stablehlo.reduce(%31093 init: %31094) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2209)
        stablehlo.return %39096 : tensor<f32> loc(#loc2209)
      } loc(#loc2209)
      %31096 = stablehlo.broadcast_in_dim %31095, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2206)
      %31097 = stablehlo.broadcast_in_dim %31096, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2210)
      %31098 = stablehlo.divide %31093, %31097 : tensor<1x1x32x2048x2048xf32> loc(#loc2210)
      %31099 = stablehlo.multiply %31096, %31096 : tensor<1x1x32x2048x1xf32> loc(#loc2211)
      %31100 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2212)
      %31101 = stablehlo.broadcast_in_dim %31100, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2213)
      %31102 = stablehlo.divide %31101, %31099 : tensor<1x1x32x2048x1xf32> loc(#loc2210)
      %31103 = stablehlo.convert %31098 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2214)
      %31104 = func.call @_einsum_654(%31103, %31060) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2215)
      %31105 = arith.select %ptrue, %31104, %31104 : tensor<1x1x2048x32x64xbf16>
      %31106 = func.call @_einsum_655(%31105, %30520#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2216)
      %31107 = arith.select %ptrue, %31106, %31106 : tensor<1x1x2048x2048xbf16>
      %31108 = stablehlo.add %31107, %30949 : tensor<1x1x2048x2048xbf16> loc(#loc2218)
      %31109 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2219)
      %31110 = stablehlo.convert %31108 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2220)
      %31111 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2221)
      %31112 = stablehlo.reduce(%31110 init: %31111) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2221)
        stablehlo.return %39096 : tensor<f32> loc(#loc2221)
      } loc(#loc2221)
      %31113 = stablehlo.broadcast_in_dim %31112, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2222)
      %31114 = stablehlo.convert %31113 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2223)
      %31115 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31116 = stablehlo.broadcast_in_dim %31115, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2224)
      %31117 = stablehlo.divide %31114, %31116 : tensor<1x1x2048x1xbf16> loc(#loc2224)
      %31118 = stablehlo.broadcast_in_dim %31117, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2225)
      %31119 = stablehlo.subtract %31108, %31118 : tensor<1x1x2048x2048xbf16> loc(#loc2225)
      %31120 = stablehlo.multiply %31119, %31119 : tensor<1x1x2048x2048xbf16> loc(#loc2226)
      %31121 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31122 = stablehlo.broadcast_in_dim %31121, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2226)
      %31123 = stablehlo.multiply %31122, %31119 : tensor<1x1x2048x2048xbf16> loc(#loc2226)
      %31124 = stablehlo.convert %31120 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2220)
      %31125 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2221)
      %31126 = stablehlo.reduce(%31124 init: %31125) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2221)
        stablehlo.return %39096 : tensor<f32> loc(#loc2221)
      } loc(#loc2221)
      %31127 = stablehlo.broadcast_in_dim %31126, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2222)
      %31128 = stablehlo.convert %31127 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2223)
      %31129 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31130 = stablehlo.broadcast_in_dim %31129, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2224)
      %31131 = stablehlo.divide %31128, %31130 : tensor<1x1x2048x1xbf16> loc(#loc2224)
      %31132 = stablehlo.broadcast_in_dim %31117, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2225)
      %31133 = stablehlo.subtract %31108, %31132 : tensor<1x1x2048x2048xbf16> loc(#loc2225)
      %31134 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31135 = stablehlo.broadcast_in_dim %31134, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2227)
      %31136 = stablehlo.add %31131, %31135 : tensor<1x1x2048x1xbf16> loc(#loc2227)
      %31137 = stablehlo.rsqrt %31136 : tensor<1x1x2048x1xbf16> loc(#loc2228)
      %31138 = stablehlo.divide %31137, %31136 : tensor<1x1x2048x1xbf16> loc(#loc2229)
      %31139 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31140 = stablehlo.broadcast_in_dim %31139, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2230)
      %31141 = stablehlo.multiply %31140, %31138 : tensor<1x1x2048x1xbf16> loc(#loc2230)
      %31142 = stablehlo.broadcast_in_dim %31137, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2230)
      %31143 = stablehlo.multiply %31133, %31142 : tensor<1x1x2048x2048xbf16> loc(#loc2230)
      %31144 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31145 = stablehlo.broadcast_in_dim %31144, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2231)
      %31146 = stablehlo.add %31145, %30515#0 : tensor<1x2048xbf16> loc(#loc2231)
      %31147 = stablehlo.broadcast_in_dim %31146, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2232)
      %31148 = stablehlo.broadcast_in_dim %31147, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2233)
      %31149 = stablehlo.multiply %31143, %31148 : tensor<1x1x2048x2048xbf16> loc(#loc2233)
      %31150 = stablehlo.broadcast_in_dim %30514#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2234)
      %31151 = stablehlo.broadcast_in_dim %31150, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2235)
      %31152 = stablehlo.add %31149, %31151 : tensor<1x1x2048x2048xbf16> loc(#loc2235)
      %31153 = func.call @_einsum_656(%31152, %30511#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2236)
      %31154 = arith.select %ptrue, %31153, %31153 : tensor<1x1x2048x8192xbf16>
      %31155 = stablehlo.broadcast_in_dim %30510#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2238)
      %31156 = stablehlo.broadcast_in_dim %31155, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2239)
      %31157 = stablehlo.add %31154, %31156 : tensor<1x1x2048x8192xbf16> loc(#loc2239)
      %31158 = stablehlo.multiply %31157, %31157 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31159 = stablehlo.multiply %31157, %31158 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31160 = stablehlo.multiply %31157, %31157 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31161 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %31162 = stablehlo.broadcast_in_dim %31161, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31163 = stablehlo.multiply %31162, %31160 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31164 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %31165 = stablehlo.broadcast_in_dim %31164, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31166 = stablehlo.multiply %31165, %31159 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31167 = stablehlo.add %31157, %31166 : tensor<1x1x2048x8192xbf16> loc(#loc2241)
      %31168 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %31169 = stablehlo.broadcast_in_dim %31168, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31170 = stablehlo.multiply %31169, %31167 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31171 = stablehlo.tanh %31170 : tensor<1x1x2048x8192xbf16> loc(#loc2242)
      %31172 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31173 = stablehlo.broadcast_in_dim %31172, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2243)
      %31174 = stablehlo.subtract %31173, %31171 : tensor<1x1x2048x8192xbf16> loc(#loc2243)
      %31175 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31176 = stablehlo.broadcast_in_dim %31175, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2241)
      %31177 = stablehlo.add %31176, %31171 : tensor<1x1x2048x8192xbf16> loc(#loc2241)
      %31178 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31179 = stablehlo.broadcast_in_dim %31178, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31180 = stablehlo.multiply %31179, %31177 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31181 = stablehlo.multiply %31157, %31180 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %31182 = stablehlo.convert %31109 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2244)
      %31183 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31184 = stablehlo.broadcast_in_dim %31183, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2245)
      %31185 = stablehlo.subtract %31184, %31182 : tensor<1x1x2048x1xf32> loc(#loc2245)
      %31186 = stablehlo.convert %31185 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2246)
      %31187 = stablehlo.broadcast_in_dim %31186, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2247)
      %31188 = stablehlo.multiply %31181, %31187 : tensor<1x1x2048x8192xbf16> loc(#loc2247)
      %31189 = func.call @_einsum_657(%31188, %30513#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2248)
      %31190 = arith.select %ptrue, %31189, %31189 : tensor<1x1x2048x2048xbf16>
      %31191 = stablehlo.broadcast_in_dim %30512#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2250)
      %31192 = stablehlo.broadcast_in_dim %31191, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2251)
      %31193 = stablehlo.add %31190, %31192 : tensor<1x1x2048x2048xbf16> loc(#loc2251)
      %31194 = stablehlo.convert %31109 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2244)
      %31195 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31196 = stablehlo.broadcast_in_dim %31195, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2245)
      %31197 = stablehlo.subtract %31196, %31194 : tensor<1x1x2048x1xf32> loc(#loc2245)
      %31198 = stablehlo.convert %31197 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2252)
      %31199 = stablehlo.broadcast_in_dim %31198, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2253)
      %31200 = stablehlo.multiply %31193, %31199 : tensor<1x1x2048x2048xbf16> loc(#loc2253)
      %31201 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31202 = stablehlo.broadcast_in_dim %31201, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2254)
      %31203 = stablehlo.multiply %31200, %31202 : tensor<1x1x2048x2048xbf16> loc(#loc2254)
      %31204 = stablehlo.add %31108, %31203 : tensor<1x1x2048x2048xbf16> loc(#loc2255)
      %31205 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2256)
      %31206 = stablehlo.convert %31204 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2257)
      %31207 = stablehlo.convert %31205 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2258)
      %31208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31209 = stablehlo.broadcast_in_dim %31208, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2259)
      %31210 = stablehlo.subtract %31209, %31207 : tensor<1x1x2048x1xf32> loc(#loc2259)
      %31211 = stablehlo.convert %31210 : tensor<1x1x2048x1xf32> loc(#loc2260)
      %31212 = stablehlo.broadcast_in_dim %31211, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31213 = stablehlo.multiply %31206, %31212 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31214 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2262)
      %31215 = stablehlo.reduce(%31213 init: %31214) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2262)
        stablehlo.return %39096 : tensor<f32> loc(#loc2262)
      } loc(#loc2262)
      %31216 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31217 = stablehlo.broadcast_in_dim %31216, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2263)
      %31218 = stablehlo.convert %31210 : tensor<1x1x2048x1xf32> loc(#loc2260)
      %31219 = stablehlo.broadcast_in_dim %31217, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2264)
      %31220 = stablehlo.broadcast_in_dim %31218, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31221 = stablehlo.multiply %31219, %31220 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31222 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2262)
      %31223 = stablehlo.reduce(%31221 init: %31222) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2262)
        stablehlo.return %39096 : tensor<f32> loc(#loc2262)
      } loc(#loc2262)
      %31224 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31225 = stablehlo.broadcast_in_dim %31224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2265)
      %31226 = stablehlo.maximum %31225, %31223 : tensor<1xf32> loc(#loc2265)
      %31227 = stablehlo.divide %31215, %31226 : tensor<1xf32> loc(#loc2266)
      %31228 = stablehlo.broadcast_in_dim %31227, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2267)
      %31229 = stablehlo.broadcast_in_dim %31228, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2268)
      %31230 = stablehlo.subtract %31206, %31229 : tensor<1x1x2048x2048xf32> loc(#loc2268)
      %31231 = stablehlo.convert %31210 : tensor<1x1x2048x1xf32> loc(#loc2260)
      %31232 = stablehlo.broadcast_in_dim %31231, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31233 = stablehlo.multiply %31230, %31232 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31234 = stablehlo.multiply %31233, %31233 : tensor<1x1x2048x2048xf32> loc(#loc2269)
      %31235 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %31236 = stablehlo.broadcast_in_dim %31235, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2269)
      %31237 = stablehlo.multiply %31236, %31233 : tensor<1x1x2048x2048xf32> loc(#loc2269)
      %31238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2262)
      %31239 = stablehlo.reduce(%31234 init: %31238) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2262)
        stablehlo.return %39096 : tensor<f32> loc(#loc2262)
      } loc(#loc2262)
      %31240 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31241 = stablehlo.broadcast_in_dim %31240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2265)
      %31242 = stablehlo.maximum %31241, %31223 : tensor<1xf32> loc(#loc2265)
      %31243 = stablehlo.divide %31239, %31242 : tensor<1xf32> loc(#loc2266)
      %31244 = stablehlo.sqrt %31243 : tensor<1xf32> loc(#loc2270)
      %31245 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31246 = stablehlo.broadcast_in_dim %31245, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2271)
      %31247 = stablehlo.divide %31246, %31244 : tensor<1xf32> loc(#loc2271)
      %31248 = stablehlo.convert %31210 : tensor<1x1x2048x1xf32> loc(#loc2260)
      %31249 = stablehlo.broadcast_in_dim %31248, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31250 = stablehlo.multiply %31206, %31249 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %31251 = stablehlo.abs %31250 : tensor<1x1x2048x2048xf32> loc(#loc2272)
      %31252 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %31253 = stablehlo.broadcast_in_dim %31252, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2273)
      %31254 = stablehlo.compare  GE, %31250, %31253,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2273)
      %31255 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2274)
      %31256 = stablehlo.reduce(%31251 init: %31255) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2274)
        stablehlo.return %39096 : tensor<f32> loc(#loc2274)
      } loc(#loc2274)
      %31257 = stablehlo.reshape %31256 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2275)
      %31258 = stablehlo.broadcast_in_dim %31257, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2276)
      %31259 = stablehlo.compare  EQ, %31251, %31258,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2276)
      %31260 = stablehlo.convert %31259 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2277)
      %31261 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2278)
      %31262 = stablehlo.reduce(%31260 init: %31261) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2278)
        stablehlo.return %39096 : tensor<f32> loc(#loc2278)
      } loc(#loc2278)
      %31263 = stablehlo.convert %31204 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2279)
      %31264 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2280)
      %31265 = stablehlo.reduce(%31263 init: %31264) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2280)
        stablehlo.return %39096 : tensor<f32> loc(#loc2280)
      } loc(#loc2280)
      %31266 = stablehlo.broadcast_in_dim %31265, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2281)
      %31267 = stablehlo.convert %31266 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2282)
      %31268 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31269 = stablehlo.broadcast_in_dim %31268, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2283)
      %31270 = stablehlo.divide %31267, %31269 : tensor<1x1x2048x1xbf16> loc(#loc2283)
      %31271 = stablehlo.broadcast_in_dim %31270, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2284)
      %31272 = stablehlo.subtract %31204, %31271 : tensor<1x1x2048x2048xbf16> loc(#loc2284)
      %31273 = stablehlo.multiply %31272, %31272 : tensor<1x1x2048x2048xbf16> loc(#loc2285)
      %31274 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31275 = stablehlo.broadcast_in_dim %31274, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2285)
      %31276 = stablehlo.multiply %31275, %31272 : tensor<1x1x2048x2048xbf16> loc(#loc2285)
      %31277 = stablehlo.convert %31273 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2279)
      %31278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2280)
      %31279 = stablehlo.reduce(%31277 init: %31278) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2280)
        stablehlo.return %39096 : tensor<f32> loc(#loc2280)
      } loc(#loc2280)
      %31280 = stablehlo.broadcast_in_dim %31279, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2281)
      %31281 = stablehlo.convert %31280 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2282)
      %31282 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31283 = stablehlo.broadcast_in_dim %31282, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2283)
      %31284 = stablehlo.divide %31281, %31283 : tensor<1x1x2048x1xbf16> loc(#loc2283)
      %31285 = stablehlo.broadcast_in_dim %31270, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2284)
      %31286 = stablehlo.subtract %31204, %31285 : tensor<1x1x2048x2048xbf16> loc(#loc2284)
      %31287 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31288 = stablehlo.broadcast_in_dim %31287, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2286)
      %31289 = stablehlo.add %31284, %31288 : tensor<1x1x2048x1xbf16> loc(#loc2286)
      %31290 = stablehlo.rsqrt %31289 : tensor<1x1x2048x1xbf16> loc(#loc2287)
      %31291 = stablehlo.divide %31290, %31289 : tensor<1x1x2048x1xbf16> loc(#loc2288)
      %31292 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31293 = stablehlo.broadcast_in_dim %31292, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2289)
      %31294 = stablehlo.multiply %31293, %31291 : tensor<1x1x2048x1xbf16> loc(#loc2289)
      %31295 = stablehlo.broadcast_in_dim %31290, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2289)
      %31296 = stablehlo.multiply %31286, %31295 : tensor<1x1x2048x2048xbf16> loc(#loc2289)
      %31297 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31298 = stablehlo.broadcast_in_dim %31297, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2290)
      %31299 = stablehlo.add %31298, %30594#0 : tensor<1x2048xbf16> loc(#loc2290)
      %31300 = stablehlo.broadcast_in_dim %31299, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2291)
      %31301 = stablehlo.broadcast_in_dim %31300, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2292)
      %31302 = stablehlo.multiply %31296, %31301 : tensor<1x1x2048x2048xbf16> loc(#loc2292)
      %31303 = stablehlo.broadcast_in_dim %30593#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2293)
      %31304 = stablehlo.broadcast_in_dim %31303, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2294)
      %31305 = stablehlo.add %31302, %31304 : tensor<1x1x2048x2048xbf16> loc(#loc2294)
      %31306 = func.call @_einsum_658(%31305, %30595#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2295)
      %31307 = "stablehlo.slice"(%31306) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2296)
      %31308 = stablehlo.reshape %31307 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2297)
      %31309 = "stablehlo.slice"(%31306) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2298)
      %31310 = stablehlo.reshape %31309 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2297)
      %31311 = "stablehlo.slice"(%31306) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2299)
      %31312 = stablehlo.reshape %31311 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2297)
      %31313 = arith.select %ptrue, %31308, %31308 : tensor<1x1x2048x32x64xbf16>
      %31314 = arith.select %ptrue, %31310, %31310 : tensor<1x1x2048x32x64xbf16>
      %31315 = arith.select %ptrue, %31312, %31312 : tensor<1x1x2048x32x64xbf16>
      %31316:2 = func.call @softplus_659(%30596#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2301)
      %31317 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %31318 = stablehlo.broadcast_in_dim %31317, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2302)
      %31319 = stablehlo.multiply %31318, %31316#0 : tensor<1x64xbf16> loc(#loc2302)
      %31320 = stablehlo.broadcast_in_dim %31319, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2303)
      %31321 = stablehlo.broadcast_in_dim %31320, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2304)
      %31322 = stablehlo.multiply %31313, %31321 : tensor<1x1x2048x32x64xbf16> loc(#loc2304)
      %31323 = func.call @_einsum_661(%31322, %31314) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2305)
      %31324 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31325 = stablehlo.broadcast_in_dim %31324, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2306)
      %31326 = stablehlo.divide %31323, %31325 : tensor<1x1x32x2048x2048xbf16> loc(#loc2306)
      %31327 = stablehlo.tanh %31326 : tensor<1x1x32x2048x2048xbf16> loc(#loc2307)
      %31328 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31329 = stablehlo.broadcast_in_dim %31328, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2308)
      %31330 = stablehlo.subtract %31329, %31327 : tensor<1x1x32x2048x2048xbf16> loc(#loc2308)
      %31331 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31332 = stablehlo.broadcast_in_dim %31331, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2309)
      %31333 = stablehlo.multiply %31332, %31327 : tensor<1x1x32x2048x2048xbf16> loc(#loc2309)
      %31334 = stablehlo.convert %31333 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2310)
      %31335 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31336 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31337 = stablehlo.multiply %31335, %31336 : tensor<f32> loc(#loc2311)
      %31338 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2312)
      %31339 = stablehlo.broadcast_in_dim %31337, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2313)
      %31340 = stablehlo.compare  GE, %31338, %31339,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2313)
      %31341 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31342:2 = func.call @_where_662(%31340, %31334, %31341) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2314)
      %31343 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2315)
      %31344 = stablehlo.reduce(%31342#0 init: %31343) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2315)
        stablehlo.return %39096 : tensor<f32> loc(#loc2315)
      } loc(#loc2315)
      %31345 = stablehlo.broadcast_in_dim %31344, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2316)
      %31346 = stablehlo.broadcast_in_dim %31345, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2317)
      %31347 = stablehlo.subtract %31342#0, %31346 : tensor<1x1x32x2048x2048xf32> loc(#loc2317)
      %31348 = stablehlo.exponential %31347 : tensor<1x1x32x2048x2048xf32> loc(#loc2318)
      %31349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2319)
      %31350 = stablehlo.reduce(%31348 init: %31349) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2319)
        stablehlo.return %39096 : tensor<f32> loc(#loc2319)
      } loc(#loc2319)
      %31351 = stablehlo.broadcast_in_dim %31350, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2316)
      %31352 = stablehlo.broadcast_in_dim %31351, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2320)
      %31353 = stablehlo.divide %31348, %31352 : tensor<1x1x32x2048x2048xf32> loc(#loc2320)
      %31354 = stablehlo.multiply %31351, %31351 : tensor<1x1x32x2048x1xf32> loc(#loc2321)
      %31355 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2322)
      %31356 = stablehlo.broadcast_in_dim %31355, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2323)
      %31357 = stablehlo.divide %31356, %31354 : tensor<1x1x32x2048x1xf32> loc(#loc2320)
      %31358 = stablehlo.convert %31353 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2324)
      %31359 = func.call @_einsum_663(%31358, %31315) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2325)
      %31360 = arith.select %ptrue, %31359, %31359 : tensor<1x1x2048x32x64xbf16>
      %31361 = func.call @_einsum_664(%31360, %30597#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2326)
      %31362 = arith.select %ptrue, %31361, %31361 : tensor<1x1x2048x2048xbf16>
      %31363 = stablehlo.add %31362, %31204 : tensor<1x1x2048x2048xbf16> loc(#loc2328)
      %31364 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2329)
      %31365 = stablehlo.convert %31363 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2330)
      %31366 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2331)
      %31367 = stablehlo.reduce(%31365 init: %31366) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2331)
        stablehlo.return %39096 : tensor<f32> loc(#loc2331)
      } loc(#loc2331)
      %31368 = stablehlo.broadcast_in_dim %31367, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2332)
      %31369 = stablehlo.convert %31368 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2333)
      %31370 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31371 = stablehlo.broadcast_in_dim %31370, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2334)
      %31372 = stablehlo.divide %31369, %31371 : tensor<1x1x2048x1xbf16> loc(#loc2334)
      %31373 = stablehlo.broadcast_in_dim %31372, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2335)
      %31374 = stablehlo.subtract %31363, %31373 : tensor<1x1x2048x2048xbf16> loc(#loc2335)
      %31375 = stablehlo.multiply %31374, %31374 : tensor<1x1x2048x2048xbf16> loc(#loc2336)
      %31376 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31377 = stablehlo.broadcast_in_dim %31376, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2336)
      %31378 = stablehlo.multiply %31377, %31374 : tensor<1x1x2048x2048xbf16> loc(#loc2336)
      %31379 = stablehlo.convert %31375 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2330)
      %31380 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2331)
      %31381 = stablehlo.reduce(%31379 init: %31380) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2331)
        stablehlo.return %39096 : tensor<f32> loc(#loc2331)
      } loc(#loc2331)
      %31382 = stablehlo.broadcast_in_dim %31381, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2332)
      %31383 = stablehlo.convert %31382 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2333)
      %31384 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31385 = stablehlo.broadcast_in_dim %31384, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2334)
      %31386 = stablehlo.divide %31383, %31385 : tensor<1x1x2048x1xbf16> loc(#loc2334)
      %31387 = stablehlo.broadcast_in_dim %31372, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2335)
      %31388 = stablehlo.subtract %31363, %31387 : tensor<1x1x2048x2048xbf16> loc(#loc2335)
      %31389 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31390 = stablehlo.broadcast_in_dim %31389, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2337)
      %31391 = stablehlo.add %31386, %31390 : tensor<1x1x2048x1xbf16> loc(#loc2337)
      %31392 = stablehlo.rsqrt %31391 : tensor<1x1x2048x1xbf16> loc(#loc2338)
      %31393 = stablehlo.divide %31392, %31391 : tensor<1x1x2048x1xbf16> loc(#loc2339)
      %31394 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31395 = stablehlo.broadcast_in_dim %31394, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2340)
      %31396 = stablehlo.multiply %31395, %31393 : tensor<1x1x2048x1xbf16> loc(#loc2340)
      %31397 = stablehlo.broadcast_in_dim %31392, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2340)
      %31398 = stablehlo.multiply %31388, %31397 : tensor<1x1x2048x2048xbf16> loc(#loc2340)
      %31399 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31400 = stablehlo.broadcast_in_dim %31399, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2341)
      %31401 = stablehlo.add %31400, %30592#0 : tensor<1x2048xbf16> loc(#loc2341)
      %31402 = stablehlo.broadcast_in_dim %31401, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2342)
      %31403 = stablehlo.broadcast_in_dim %31402, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2343)
      %31404 = stablehlo.multiply %31398, %31403 : tensor<1x1x2048x2048xbf16> loc(#loc2343)
      %31405 = stablehlo.broadcast_in_dim %30591#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2344)
      %31406 = stablehlo.broadcast_in_dim %31405, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2345)
      %31407 = stablehlo.add %31404, %31406 : tensor<1x1x2048x2048xbf16> loc(#loc2345)
      %31408 = func.call @_einsum_665(%31407, %30588#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2346)
      %31409 = arith.select %ptrue, %31408, %31408 : tensor<1x1x2048x8192xbf16>
      %31410 = stablehlo.broadcast_in_dim %30587#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2348)
      %31411 = stablehlo.broadcast_in_dim %31410, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2349)
      %31412 = stablehlo.add %31409, %31411 : tensor<1x1x2048x8192xbf16> loc(#loc2349)
      %31413 = stablehlo.multiply %31412, %31412 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31414 = stablehlo.multiply %31412, %31413 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31415 = stablehlo.multiply %31412, %31412 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31416 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %31417 = stablehlo.broadcast_in_dim %31416, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31418 = stablehlo.multiply %31417, %31415 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31419 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %31420 = stablehlo.broadcast_in_dim %31419, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31421 = stablehlo.multiply %31420, %31414 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31422 = stablehlo.add %31412, %31421 : tensor<1x1x2048x8192xbf16> loc(#loc2351)
      %31423 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %31424 = stablehlo.broadcast_in_dim %31423, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31425 = stablehlo.multiply %31424, %31422 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31426 = stablehlo.tanh %31425 : tensor<1x1x2048x8192xbf16> loc(#loc2352)
      %31427 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31428 = stablehlo.broadcast_in_dim %31427, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2353)
      %31429 = stablehlo.subtract %31428, %31426 : tensor<1x1x2048x8192xbf16> loc(#loc2353)
      %31430 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31431 = stablehlo.broadcast_in_dim %31430, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2351)
      %31432 = stablehlo.add %31431, %31426 : tensor<1x1x2048x8192xbf16> loc(#loc2351)
      %31433 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31434 = stablehlo.broadcast_in_dim %31433, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31435 = stablehlo.multiply %31434, %31432 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31436 = stablehlo.multiply %31412, %31435 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %31437 = stablehlo.convert %31364 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2354)
      %31438 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31439 = stablehlo.broadcast_in_dim %31438, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2355)
      %31440 = stablehlo.subtract %31439, %31437 : tensor<1x1x2048x1xf32> loc(#loc2355)
      %31441 = stablehlo.convert %31440 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2356)
      %31442 = stablehlo.broadcast_in_dim %31441, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2357)
      %31443 = stablehlo.multiply %31436, %31442 : tensor<1x1x2048x8192xbf16> loc(#loc2357)
      %31444 = func.call @_einsum_666(%31443, %30590#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2358)
      %31445 = arith.select %ptrue, %31444, %31444 : tensor<1x1x2048x2048xbf16>
      %31446 = stablehlo.broadcast_in_dim %30589#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2360)
      %31447 = stablehlo.broadcast_in_dim %31446, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2361)
      %31448 = stablehlo.add %31445, %31447 : tensor<1x1x2048x2048xbf16> loc(#loc2361)
      %31449 = stablehlo.convert %31364 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2354)
      %31450 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31451 = stablehlo.broadcast_in_dim %31450, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2355)
      %31452 = stablehlo.subtract %31451, %31449 : tensor<1x1x2048x1xf32> loc(#loc2355)
      %31453 = stablehlo.convert %31452 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2362)
      %31454 = stablehlo.broadcast_in_dim %31453, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2363)
      %31455 = stablehlo.multiply %31448, %31454 : tensor<1x1x2048x2048xbf16> loc(#loc2363)
      %31456 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31457 = stablehlo.broadcast_in_dim %31456, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2364)
      %31458 = stablehlo.multiply %31455, %31457 : tensor<1x1x2048x2048xbf16> loc(#loc2364)
      %31459 = stablehlo.add %31363, %31458 : tensor<1x1x2048x2048xbf16> loc(#loc2365)
      %31460 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2366)
      %31461 = stablehlo.convert %31459 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2367)
      %31462 = stablehlo.convert %31460 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2368)
      %31463 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31464 = stablehlo.broadcast_in_dim %31463, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2369)
      %31465 = stablehlo.subtract %31464, %31462 : tensor<1x1x2048x1xf32> loc(#loc2369)
      %31466 = stablehlo.convert %31465 : tensor<1x1x2048x1xf32> loc(#loc2370)
      %31467 = stablehlo.broadcast_in_dim %31466, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31468 = stablehlo.multiply %31461, %31467 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31469 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2372)
      %31470 = stablehlo.reduce(%31468 init: %31469) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2372)
        stablehlo.return %39096 : tensor<f32> loc(#loc2372)
      } loc(#loc2372)
      %31471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31472 = stablehlo.broadcast_in_dim %31471, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2373)
      %31473 = stablehlo.convert %31465 : tensor<1x1x2048x1xf32> loc(#loc2370)
      %31474 = stablehlo.broadcast_in_dim %31472, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2374)
      %31475 = stablehlo.broadcast_in_dim %31473, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31476 = stablehlo.multiply %31474, %31475 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31477 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2372)
      %31478 = stablehlo.reduce(%31476 init: %31477) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2372)
        stablehlo.return %39096 : tensor<f32> loc(#loc2372)
      } loc(#loc2372)
      %31479 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31480 = stablehlo.broadcast_in_dim %31479, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2375)
      %31481 = stablehlo.maximum %31480, %31478 : tensor<1xf32> loc(#loc2375)
      %31482 = stablehlo.divide %31470, %31481 : tensor<1xf32> loc(#loc2376)
      %31483 = stablehlo.broadcast_in_dim %31482, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2377)
      %31484 = stablehlo.broadcast_in_dim %31483, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2378)
      %31485 = stablehlo.subtract %31461, %31484 : tensor<1x1x2048x2048xf32> loc(#loc2378)
      %31486 = stablehlo.convert %31465 : tensor<1x1x2048x1xf32> loc(#loc2370)
      %31487 = stablehlo.broadcast_in_dim %31486, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31488 = stablehlo.multiply %31485, %31487 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31489 = stablehlo.multiply %31488, %31488 : tensor<1x1x2048x2048xf32> loc(#loc2379)
      %31490 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %31491 = stablehlo.broadcast_in_dim %31490, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2379)
      %31492 = stablehlo.multiply %31491, %31488 : tensor<1x1x2048x2048xf32> loc(#loc2379)
      %31493 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2372)
      %31494 = stablehlo.reduce(%31489 init: %31493) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2372)
        stablehlo.return %39096 : tensor<f32> loc(#loc2372)
      } loc(#loc2372)
      %31495 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31496 = stablehlo.broadcast_in_dim %31495, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2375)
      %31497 = stablehlo.maximum %31496, %31478 : tensor<1xf32> loc(#loc2375)
      %31498 = stablehlo.divide %31494, %31497 : tensor<1xf32> loc(#loc2376)
      %31499 = stablehlo.sqrt %31498 : tensor<1xf32> loc(#loc2380)
      %31500 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31501 = stablehlo.broadcast_in_dim %31500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2381)
      %31502 = stablehlo.divide %31501, %31499 : tensor<1xf32> loc(#loc2381)
      %31503 = stablehlo.convert %31465 : tensor<1x1x2048x1xf32> loc(#loc2370)
      %31504 = stablehlo.broadcast_in_dim %31503, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31505 = stablehlo.multiply %31461, %31504 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %31506 = stablehlo.abs %31505 : tensor<1x1x2048x2048xf32> loc(#loc2382)
      %31507 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %31508 = stablehlo.broadcast_in_dim %31507, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2383)
      %31509 = stablehlo.compare  GE, %31505, %31508,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2383)
      %31510 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2384)
      %31511 = stablehlo.reduce(%31506 init: %31510) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2384)
        stablehlo.return %39096 : tensor<f32> loc(#loc2384)
      } loc(#loc2384)
      %31512 = stablehlo.reshape %31511 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2385)
      %31513 = stablehlo.broadcast_in_dim %31512, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2386)
      %31514 = stablehlo.compare  EQ, %31506, %31513,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2386)
      %31515 = stablehlo.convert %31514 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2387)
      %31516 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2388)
      %31517 = stablehlo.reduce(%31515 init: %31516) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2388)
        stablehlo.return %39096 : tensor<f32> loc(#loc2388)
      } loc(#loc2388)
      %31518 = stablehlo.convert %31459 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2389)
      %31519 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2390)
      %31520 = stablehlo.reduce(%31518 init: %31519) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2390)
        stablehlo.return %39096 : tensor<f32> loc(#loc2390)
      } loc(#loc2390)
      %31521 = stablehlo.broadcast_in_dim %31520, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2391)
      %31522 = stablehlo.convert %31521 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2392)
      %31523 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31524 = stablehlo.broadcast_in_dim %31523, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2393)
      %31525 = stablehlo.divide %31522, %31524 : tensor<1x1x2048x1xbf16> loc(#loc2393)
      %31526 = stablehlo.broadcast_in_dim %31525, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2394)
      %31527 = stablehlo.subtract %31459, %31526 : tensor<1x1x2048x2048xbf16> loc(#loc2394)
      %31528 = stablehlo.multiply %31527, %31527 : tensor<1x1x2048x2048xbf16> loc(#loc2395)
      %31529 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31530 = stablehlo.broadcast_in_dim %31529, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2395)
      %31531 = stablehlo.multiply %31530, %31527 : tensor<1x1x2048x2048xbf16> loc(#loc2395)
      %31532 = stablehlo.convert %31528 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2389)
      %31533 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2390)
      %31534 = stablehlo.reduce(%31532 init: %31533) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2390)
        stablehlo.return %39096 : tensor<f32> loc(#loc2390)
      } loc(#loc2390)
      %31535 = stablehlo.broadcast_in_dim %31534, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2391)
      %31536 = stablehlo.convert %31535 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2392)
      %31537 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31538 = stablehlo.broadcast_in_dim %31537, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2393)
      %31539 = stablehlo.divide %31536, %31538 : tensor<1x1x2048x1xbf16> loc(#loc2393)
      %31540 = stablehlo.broadcast_in_dim %31525, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2394)
      %31541 = stablehlo.subtract %31459, %31540 : tensor<1x1x2048x2048xbf16> loc(#loc2394)
      %31542 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31543 = stablehlo.broadcast_in_dim %31542, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2396)
      %31544 = stablehlo.add %31539, %31543 : tensor<1x1x2048x1xbf16> loc(#loc2396)
      %31545 = stablehlo.rsqrt %31544 : tensor<1x1x2048x1xbf16> loc(#loc2397)
      %31546 = stablehlo.divide %31545, %31544 : tensor<1x1x2048x1xbf16> loc(#loc2398)
      %31547 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31548 = stablehlo.broadcast_in_dim %31547, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2399)
      %31549 = stablehlo.multiply %31548, %31546 : tensor<1x1x2048x1xbf16> loc(#loc2399)
      %31550 = stablehlo.broadcast_in_dim %31545, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2399)
      %31551 = stablehlo.multiply %31541, %31550 : tensor<1x1x2048x2048xbf16> loc(#loc2399)
      %31552 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31553 = stablehlo.broadcast_in_dim %31552, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2400)
      %31554 = stablehlo.add %31553, %30605#0 : tensor<1x2048xbf16> loc(#loc2400)
      %31555 = stablehlo.broadcast_in_dim %31554, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2401)
      %31556 = stablehlo.broadcast_in_dim %31555, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2402)
      %31557 = stablehlo.multiply %31551, %31556 : tensor<1x1x2048x2048xbf16> loc(#loc2402)
      %31558 = stablehlo.broadcast_in_dim %30604#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2403)
      %31559 = stablehlo.broadcast_in_dim %31558, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2404)
      %31560 = stablehlo.add %31557, %31559 : tensor<1x1x2048x2048xbf16> loc(#loc2404)
      %31561 = func.call @_einsum_667(%31560, %30606#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2405)
      %31562 = "stablehlo.slice"(%31561) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2406)
      %31563 = stablehlo.reshape %31562 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2407)
      %31564 = "stablehlo.slice"(%31561) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2408)
      %31565 = stablehlo.reshape %31564 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2407)
      %31566 = "stablehlo.slice"(%31561) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2409)
      %31567 = stablehlo.reshape %31566 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2407)
      %31568 = arith.select %ptrue, %31563, %31563 : tensor<1x1x2048x32x64xbf16>
      %31569 = arith.select %ptrue, %31565, %31565 : tensor<1x1x2048x32x64xbf16>
      %31570 = arith.select %ptrue, %31567, %31567 : tensor<1x1x2048x32x64xbf16>
      %31571:2 = func.call @softplus_668(%30607#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2411)
      %31572 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %31573 = stablehlo.broadcast_in_dim %31572, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2412)
      %31574 = stablehlo.multiply %31573, %31571#0 : tensor<1x64xbf16> loc(#loc2412)
      %31575 = stablehlo.broadcast_in_dim %31574, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2413)
      %31576 = stablehlo.broadcast_in_dim %31575, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2414)
      %31577 = stablehlo.multiply %31568, %31576 : tensor<1x1x2048x32x64xbf16> loc(#loc2414)
      %31578 = func.call @_einsum_670(%31577, %31569) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2415)
      %31579 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31580 = stablehlo.broadcast_in_dim %31579, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2416)
      %31581 = stablehlo.divide %31578, %31580 : tensor<1x1x32x2048x2048xbf16> loc(#loc2416)
      %31582 = stablehlo.tanh %31581 : tensor<1x1x32x2048x2048xbf16> loc(#loc2417)
      %31583 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31584 = stablehlo.broadcast_in_dim %31583, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2418)
      %31585 = stablehlo.subtract %31584, %31582 : tensor<1x1x32x2048x2048xbf16> loc(#loc2418)
      %31586 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31587 = stablehlo.broadcast_in_dim %31586, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2419)
      %31588 = stablehlo.multiply %31587, %31582 : tensor<1x1x32x2048x2048xbf16> loc(#loc2419)
      %31589 = stablehlo.convert %31588 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2420)
      %31590 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31591 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31592 = stablehlo.multiply %31590, %31591 : tensor<f32> loc(#loc2421)
      %31593 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2422)
      %31594 = stablehlo.broadcast_in_dim %31592, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2423)
      %31595 = stablehlo.compare  GE, %31593, %31594,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2423)
      %31596 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31597:2 = func.call @_where_671(%31595, %31589, %31596) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2424)
      %31598 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2425)
      %31599 = stablehlo.reduce(%31597#0 init: %31598) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2425)
        stablehlo.return %39096 : tensor<f32> loc(#loc2425)
      } loc(#loc2425)
      %31600 = stablehlo.broadcast_in_dim %31599, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2426)
      %31601 = stablehlo.broadcast_in_dim %31600, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2427)
      %31602 = stablehlo.subtract %31597#0, %31601 : tensor<1x1x32x2048x2048xf32> loc(#loc2427)
      %31603 = stablehlo.exponential %31602 : tensor<1x1x32x2048x2048xf32> loc(#loc2428)
      %31604 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2429)
      %31605 = stablehlo.reduce(%31603 init: %31604) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2429)
        stablehlo.return %39096 : tensor<f32> loc(#loc2429)
      } loc(#loc2429)
      %31606 = stablehlo.broadcast_in_dim %31605, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2426)
      %31607 = stablehlo.broadcast_in_dim %31606, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2430)
      %31608 = stablehlo.divide %31603, %31607 : tensor<1x1x32x2048x2048xf32> loc(#loc2430)
      %31609 = stablehlo.multiply %31606, %31606 : tensor<1x1x32x2048x1xf32> loc(#loc2431)
      %31610 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2432)
      %31611 = stablehlo.broadcast_in_dim %31610, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2433)
      %31612 = stablehlo.divide %31611, %31609 : tensor<1x1x32x2048x1xf32> loc(#loc2430)
      %31613 = stablehlo.convert %31608 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2434)
      %31614 = func.call @_einsum_672(%31613, %31570) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2435)
      %31615 = arith.select %ptrue, %31614, %31614 : tensor<1x1x2048x32x64xbf16>
      %31616 = func.call @_einsum_673(%31615, %30608#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2436)
      %31617 = arith.select %ptrue, %31616, %31616 : tensor<1x1x2048x2048xbf16>
      %31618 = stablehlo.add %31617, %31459 : tensor<1x1x2048x2048xbf16> loc(#loc2438)
      %31619 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2439)
      %31620 = stablehlo.convert %31618 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2440)
      %31621 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2441)
      %31622 = stablehlo.reduce(%31620 init: %31621) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2441)
        stablehlo.return %39096 : tensor<f32> loc(#loc2441)
      } loc(#loc2441)
      %31623 = stablehlo.broadcast_in_dim %31622, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2442)
      %31624 = stablehlo.convert %31623 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2443)
      %31625 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31626 = stablehlo.broadcast_in_dim %31625, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2444)
      %31627 = stablehlo.divide %31624, %31626 : tensor<1x1x2048x1xbf16> loc(#loc2444)
      %31628 = stablehlo.broadcast_in_dim %31627, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2445)
      %31629 = stablehlo.subtract %31618, %31628 : tensor<1x1x2048x2048xbf16> loc(#loc2445)
      %31630 = stablehlo.multiply %31629, %31629 : tensor<1x1x2048x2048xbf16> loc(#loc2446)
      %31631 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31632 = stablehlo.broadcast_in_dim %31631, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2446)
      %31633 = stablehlo.multiply %31632, %31629 : tensor<1x1x2048x2048xbf16> loc(#loc2446)
      %31634 = stablehlo.convert %31630 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2440)
      %31635 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2441)
      %31636 = stablehlo.reduce(%31634 init: %31635) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2441)
        stablehlo.return %39096 : tensor<f32> loc(#loc2441)
      } loc(#loc2441)
      %31637 = stablehlo.broadcast_in_dim %31636, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2442)
      %31638 = stablehlo.convert %31637 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2443)
      %31639 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31640 = stablehlo.broadcast_in_dim %31639, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2444)
      %31641 = stablehlo.divide %31638, %31640 : tensor<1x1x2048x1xbf16> loc(#loc2444)
      %31642 = stablehlo.broadcast_in_dim %31627, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2445)
      %31643 = stablehlo.subtract %31618, %31642 : tensor<1x1x2048x2048xbf16> loc(#loc2445)
      %31644 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31645 = stablehlo.broadcast_in_dim %31644, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2447)
      %31646 = stablehlo.add %31641, %31645 : tensor<1x1x2048x1xbf16> loc(#loc2447)
      %31647 = stablehlo.rsqrt %31646 : tensor<1x1x2048x1xbf16> loc(#loc2448)
      %31648 = stablehlo.divide %31647, %31646 : tensor<1x1x2048x1xbf16> loc(#loc2449)
      %31649 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31650 = stablehlo.broadcast_in_dim %31649, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2450)
      %31651 = stablehlo.multiply %31650, %31648 : tensor<1x1x2048x1xbf16> loc(#loc2450)
      %31652 = stablehlo.broadcast_in_dim %31647, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2450)
      %31653 = stablehlo.multiply %31643, %31652 : tensor<1x1x2048x2048xbf16> loc(#loc2450)
      %31654 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31655 = stablehlo.broadcast_in_dim %31654, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2451)
      %31656 = stablehlo.add %31655, %30603#0 : tensor<1x2048xbf16> loc(#loc2451)
      %31657 = stablehlo.broadcast_in_dim %31656, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2452)
      %31658 = stablehlo.broadcast_in_dim %31657, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2453)
      %31659 = stablehlo.multiply %31653, %31658 : tensor<1x1x2048x2048xbf16> loc(#loc2453)
      %31660 = stablehlo.broadcast_in_dim %30602#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2454)
      %31661 = stablehlo.broadcast_in_dim %31660, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2455)
      %31662 = stablehlo.add %31659, %31661 : tensor<1x1x2048x2048xbf16> loc(#loc2455)
      %31663 = func.call @_einsum_674(%31662, %30599#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2456)
      %31664 = arith.select %ptrue, %31663, %31663 : tensor<1x1x2048x8192xbf16>
      %31665 = stablehlo.broadcast_in_dim %30598#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2458)
      %31666 = stablehlo.broadcast_in_dim %31665, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2459)
      %31667 = stablehlo.add %31664, %31666 : tensor<1x1x2048x8192xbf16> loc(#loc2459)
      %31668 = stablehlo.multiply %31667, %31667 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31669 = stablehlo.multiply %31667, %31668 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31670 = stablehlo.multiply %31667, %31667 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31671 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %31672 = stablehlo.broadcast_in_dim %31671, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31673 = stablehlo.multiply %31672, %31670 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31674 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %31675 = stablehlo.broadcast_in_dim %31674, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31676 = stablehlo.multiply %31675, %31669 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31677 = stablehlo.add %31667, %31676 : tensor<1x1x2048x8192xbf16> loc(#loc2461)
      %31678 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %31679 = stablehlo.broadcast_in_dim %31678, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31680 = stablehlo.multiply %31679, %31677 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31681 = stablehlo.tanh %31680 : tensor<1x1x2048x8192xbf16> loc(#loc2462)
      %31682 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31683 = stablehlo.broadcast_in_dim %31682, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2463)
      %31684 = stablehlo.subtract %31683, %31681 : tensor<1x1x2048x8192xbf16> loc(#loc2463)
      %31685 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31686 = stablehlo.broadcast_in_dim %31685, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2461)
      %31687 = stablehlo.add %31686, %31681 : tensor<1x1x2048x8192xbf16> loc(#loc2461)
      %31688 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31689 = stablehlo.broadcast_in_dim %31688, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31690 = stablehlo.multiply %31689, %31687 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31691 = stablehlo.multiply %31667, %31690 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %31692 = stablehlo.convert %31619 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2464)
      %31693 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31694 = stablehlo.broadcast_in_dim %31693, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2465)
      %31695 = stablehlo.subtract %31694, %31692 : tensor<1x1x2048x1xf32> loc(#loc2465)
      %31696 = stablehlo.convert %31695 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2466)
      %31697 = stablehlo.broadcast_in_dim %31696, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2467)
      %31698 = stablehlo.multiply %31691, %31697 : tensor<1x1x2048x8192xbf16> loc(#loc2467)
      %31699 = func.call @_einsum_675(%31698, %30601#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2468)
      %31700 = arith.select %ptrue, %31699, %31699 : tensor<1x1x2048x2048xbf16>
      %31701 = stablehlo.broadcast_in_dim %30600#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2470)
      %31702 = stablehlo.broadcast_in_dim %31701, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2471)
      %31703 = stablehlo.add %31700, %31702 : tensor<1x1x2048x2048xbf16> loc(#loc2471)
      %31704 = stablehlo.convert %31619 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2464)
      %31705 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31706 = stablehlo.broadcast_in_dim %31705, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2465)
      %31707 = stablehlo.subtract %31706, %31704 : tensor<1x1x2048x1xf32> loc(#loc2465)
      %31708 = stablehlo.convert %31707 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2472)
      %31709 = stablehlo.broadcast_in_dim %31708, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2473)
      %31710 = stablehlo.multiply %31703, %31709 : tensor<1x1x2048x2048xbf16> loc(#loc2473)
      %31711 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31712 = stablehlo.broadcast_in_dim %31711, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2474)
      %31713 = stablehlo.multiply %31710, %31712 : tensor<1x1x2048x2048xbf16> loc(#loc2474)
      %31714 = stablehlo.add %31618, %31713 : tensor<1x1x2048x2048xbf16> loc(#loc2475)
      %31715 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2476)
      %31716 = stablehlo.convert %31714 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2477)
      %31717 = stablehlo.convert %31715 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2478)
      %31718 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31719 = stablehlo.broadcast_in_dim %31718, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2479)
      %31720 = stablehlo.subtract %31719, %31717 : tensor<1x1x2048x1xf32> loc(#loc2479)
      %31721 = stablehlo.convert %31720 : tensor<1x1x2048x1xf32> loc(#loc2480)
      %31722 = stablehlo.broadcast_in_dim %31721, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31723 = stablehlo.multiply %31716, %31722 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31724 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2482)
      %31725 = stablehlo.reduce(%31723 init: %31724) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2482)
        stablehlo.return %39096 : tensor<f32> loc(#loc2482)
      } loc(#loc2482)
      %31726 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31727 = stablehlo.broadcast_in_dim %31726, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2483)
      %31728 = stablehlo.convert %31720 : tensor<1x1x2048x1xf32> loc(#loc2480)
      %31729 = stablehlo.broadcast_in_dim %31727, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2484)
      %31730 = stablehlo.broadcast_in_dim %31728, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31731 = stablehlo.multiply %31729, %31730 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31732 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2482)
      %31733 = stablehlo.reduce(%31731 init: %31732) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2482)
        stablehlo.return %39096 : tensor<f32> loc(#loc2482)
      } loc(#loc2482)
      %31734 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31735 = stablehlo.broadcast_in_dim %31734, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2485)
      %31736 = stablehlo.maximum %31735, %31733 : tensor<1xf32> loc(#loc2485)
      %31737 = stablehlo.divide %31725, %31736 : tensor<1xf32> loc(#loc2486)
      %31738 = stablehlo.broadcast_in_dim %31737, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2487)
      %31739 = stablehlo.broadcast_in_dim %31738, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2488)
      %31740 = stablehlo.subtract %31716, %31739 : tensor<1x1x2048x2048xf32> loc(#loc2488)
      %31741 = stablehlo.convert %31720 : tensor<1x1x2048x1xf32> loc(#loc2480)
      %31742 = stablehlo.broadcast_in_dim %31741, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31743 = stablehlo.multiply %31740, %31742 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31744 = stablehlo.multiply %31743, %31743 : tensor<1x1x2048x2048xf32> loc(#loc2489)
      %31745 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %31746 = stablehlo.broadcast_in_dim %31745, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2489)
      %31747 = stablehlo.multiply %31746, %31743 : tensor<1x1x2048x2048xf32> loc(#loc2489)
      %31748 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2482)
      %31749 = stablehlo.reduce(%31744 init: %31748) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2482)
        stablehlo.return %39096 : tensor<f32> loc(#loc2482)
      } loc(#loc2482)
      %31750 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31751 = stablehlo.broadcast_in_dim %31750, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2485)
      %31752 = stablehlo.maximum %31751, %31733 : tensor<1xf32> loc(#loc2485)
      %31753 = stablehlo.divide %31749, %31752 : tensor<1xf32> loc(#loc2486)
      %31754 = stablehlo.sqrt %31753 : tensor<1xf32> loc(#loc2490)
      %31755 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31756 = stablehlo.broadcast_in_dim %31755, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2491)
      %31757 = stablehlo.divide %31756, %31754 : tensor<1xf32> loc(#loc2491)
      %31758 = stablehlo.convert %31720 : tensor<1x1x2048x1xf32> loc(#loc2480)
      %31759 = stablehlo.broadcast_in_dim %31758, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31760 = stablehlo.multiply %31716, %31759 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %31761 = stablehlo.abs %31760 : tensor<1x1x2048x2048xf32> loc(#loc2492)
      %31762 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %31763 = stablehlo.broadcast_in_dim %31762, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2493)
      %31764 = stablehlo.compare  GE, %31760, %31763,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2493)
      %31765 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2494)
      %31766 = stablehlo.reduce(%31761 init: %31765) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2494)
        stablehlo.return %39096 : tensor<f32> loc(#loc2494)
      } loc(#loc2494)
      %31767 = stablehlo.reshape %31766 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2495)
      %31768 = stablehlo.broadcast_in_dim %31767, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2496)
      %31769 = stablehlo.compare  EQ, %31761, %31768,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2496)
      %31770 = stablehlo.convert %31769 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2497)
      %31771 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2498)
      %31772 = stablehlo.reduce(%31770 init: %31771) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2498)
        stablehlo.return %39096 : tensor<f32> loc(#loc2498)
      } loc(#loc2498)
      %31773 = stablehlo.convert %31714 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2499)
      %31774 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2500)
      %31775 = stablehlo.reduce(%31773 init: %31774) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2500)
        stablehlo.return %39096 : tensor<f32> loc(#loc2500)
      } loc(#loc2500)
      %31776 = stablehlo.broadcast_in_dim %31775, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2501)
      %31777 = stablehlo.convert %31776 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2502)
      %31778 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31779 = stablehlo.broadcast_in_dim %31778, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2503)
      %31780 = stablehlo.divide %31777, %31779 : tensor<1x1x2048x1xbf16> loc(#loc2503)
      %31781 = stablehlo.broadcast_in_dim %31780, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2504)
      %31782 = stablehlo.subtract %31714, %31781 : tensor<1x1x2048x2048xbf16> loc(#loc2504)
      %31783 = stablehlo.multiply %31782, %31782 : tensor<1x1x2048x2048xbf16> loc(#loc2505)
      %31784 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31785 = stablehlo.broadcast_in_dim %31784, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2505)
      %31786 = stablehlo.multiply %31785, %31782 : tensor<1x1x2048x2048xbf16> loc(#loc2505)
      %31787 = stablehlo.convert %31783 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2499)
      %31788 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2500)
      %31789 = stablehlo.reduce(%31787 init: %31788) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2500)
        stablehlo.return %39096 : tensor<f32> loc(#loc2500)
      } loc(#loc2500)
      %31790 = stablehlo.broadcast_in_dim %31789, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2501)
      %31791 = stablehlo.convert %31790 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2502)
      %31792 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31793 = stablehlo.broadcast_in_dim %31792, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2503)
      %31794 = stablehlo.divide %31791, %31793 : tensor<1x1x2048x1xbf16> loc(#loc2503)
      %31795 = stablehlo.broadcast_in_dim %31780, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2504)
      %31796 = stablehlo.subtract %31714, %31795 : tensor<1x1x2048x2048xbf16> loc(#loc2504)
      %31797 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31798 = stablehlo.broadcast_in_dim %31797, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2506)
      %31799 = stablehlo.add %31794, %31798 : tensor<1x1x2048x1xbf16> loc(#loc2506)
      %31800 = stablehlo.rsqrt %31799 : tensor<1x1x2048x1xbf16> loc(#loc2507)
      %31801 = stablehlo.divide %31800, %31799 : tensor<1x1x2048x1xbf16> loc(#loc2508)
      %31802 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31803 = stablehlo.broadcast_in_dim %31802, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2509)
      %31804 = stablehlo.multiply %31803, %31801 : tensor<1x1x2048x1xbf16> loc(#loc2509)
      %31805 = stablehlo.broadcast_in_dim %31800, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2509)
      %31806 = stablehlo.multiply %31796, %31805 : tensor<1x1x2048x2048xbf16> loc(#loc2509)
      %31807 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31808 = stablehlo.broadcast_in_dim %31807, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2510)
      %31809 = stablehlo.add %31808, %30616#0 : tensor<1x2048xbf16> loc(#loc2510)
      %31810 = stablehlo.broadcast_in_dim %31809, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2511)
      %31811 = stablehlo.broadcast_in_dim %31810, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2512)
      %31812 = stablehlo.multiply %31806, %31811 : tensor<1x1x2048x2048xbf16> loc(#loc2512)
      %31813 = stablehlo.broadcast_in_dim %30615#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2513)
      %31814 = stablehlo.broadcast_in_dim %31813, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2514)
      %31815 = stablehlo.add %31812, %31814 : tensor<1x1x2048x2048xbf16> loc(#loc2514)
      %31816 = func.call @_einsum_676(%31815, %30617#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2515)
      %31817 = "stablehlo.slice"(%31816) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2516)
      %31818 = stablehlo.reshape %31817 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2517)
      %31819 = "stablehlo.slice"(%31816) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2518)
      %31820 = stablehlo.reshape %31819 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2517)
      %31821 = "stablehlo.slice"(%31816) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2519)
      %31822 = stablehlo.reshape %31821 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2517)
      %31823 = arith.select %ptrue, %31818, %31818 : tensor<1x1x2048x32x64xbf16>
      %31824 = arith.select %ptrue, %31820, %31820 : tensor<1x1x2048x32x64xbf16>
      %31825 = arith.select %ptrue, %31822, %31822 : tensor<1x1x2048x32x64xbf16>
      %31826:2 = func.call @softplus_677(%30618#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2521)
      %31827 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %31828 = stablehlo.broadcast_in_dim %31827, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2522)
      %31829 = stablehlo.multiply %31828, %31826#0 : tensor<1x64xbf16> loc(#loc2522)
      %31830 = stablehlo.broadcast_in_dim %31829, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2523)
      %31831 = stablehlo.broadcast_in_dim %31830, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2524)
      %31832 = stablehlo.multiply %31823, %31831 : tensor<1x1x2048x32x64xbf16> loc(#loc2524)
      %31833 = func.call @_einsum_679(%31832, %31824) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2525)
      %31834 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31835 = stablehlo.broadcast_in_dim %31834, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2526)
      %31836 = stablehlo.divide %31833, %31835 : tensor<1x1x32x2048x2048xbf16> loc(#loc2526)
      %31837 = stablehlo.tanh %31836 : tensor<1x1x32x2048x2048xbf16> loc(#loc2527)
      %31838 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31839 = stablehlo.broadcast_in_dim %31838, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2528)
      %31840 = stablehlo.subtract %31839, %31837 : tensor<1x1x32x2048x2048xbf16> loc(#loc2528)
      %31841 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %31842 = stablehlo.broadcast_in_dim %31841, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2529)
      %31843 = stablehlo.multiply %31842, %31837 : tensor<1x1x32x2048x2048xbf16> loc(#loc2529)
      %31844 = stablehlo.convert %31843 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2530)
      %31845 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31846 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %31847 = stablehlo.multiply %31845, %31846 : tensor<f32> loc(#loc2531)
      %31848 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2532)
      %31849 = stablehlo.broadcast_in_dim %31847, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2533)
      %31850 = stablehlo.compare  GE, %31848, %31849,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2533)
      %31851 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %31852:2 = func.call @_where_680(%31850, %31844, %31851) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2534)
      %31853 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2535)
      %31854 = stablehlo.reduce(%31852#0 init: %31853) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2535)
        stablehlo.return %39096 : tensor<f32> loc(#loc2535)
      } loc(#loc2535)
      %31855 = stablehlo.broadcast_in_dim %31854, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2536)
      %31856 = stablehlo.broadcast_in_dim %31855, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2537)
      %31857 = stablehlo.subtract %31852#0, %31856 : tensor<1x1x32x2048x2048xf32> loc(#loc2537)
      %31858 = stablehlo.exponential %31857 : tensor<1x1x32x2048x2048xf32> loc(#loc2538)
      %31859 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2539)
      %31860 = stablehlo.reduce(%31858 init: %31859) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2539)
        stablehlo.return %39096 : tensor<f32> loc(#loc2539)
      } loc(#loc2539)
      %31861 = stablehlo.broadcast_in_dim %31860, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2536)
      %31862 = stablehlo.broadcast_in_dim %31861, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2540)
      %31863 = stablehlo.divide %31858, %31862 : tensor<1x1x32x2048x2048xf32> loc(#loc2540)
      %31864 = stablehlo.multiply %31861, %31861 : tensor<1x1x32x2048x1xf32> loc(#loc2541)
      %31865 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2542)
      %31866 = stablehlo.broadcast_in_dim %31865, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2543)
      %31867 = stablehlo.divide %31866, %31864 : tensor<1x1x32x2048x1xf32> loc(#loc2540)
      %31868 = stablehlo.convert %31863 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2544)
      %31869 = func.call @_einsum_681(%31868, %31825) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2545)
      %31870 = arith.select %ptrue, %31869, %31869 : tensor<1x1x2048x32x64xbf16>
      %31871 = func.call @_einsum_682(%31870, %30619#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2546)
      %31872 = arith.select %ptrue, %31871, %31871 : tensor<1x1x2048x2048xbf16>
      %31873 = stablehlo.add %31872, %31714 : tensor<1x1x2048x2048xbf16> loc(#loc2548)
      %31874 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2549)
      %31875 = stablehlo.convert %31873 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2550)
      %31876 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2551)
      %31877 = stablehlo.reduce(%31875 init: %31876) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2551)
        stablehlo.return %39096 : tensor<f32> loc(#loc2551)
      } loc(#loc2551)
      %31878 = stablehlo.broadcast_in_dim %31877, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2552)
      %31879 = stablehlo.convert %31878 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2553)
      %31880 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31881 = stablehlo.broadcast_in_dim %31880, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2554)
      %31882 = stablehlo.divide %31879, %31881 : tensor<1x1x2048x1xbf16> loc(#loc2554)
      %31883 = stablehlo.broadcast_in_dim %31882, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2555)
      %31884 = stablehlo.subtract %31873, %31883 : tensor<1x1x2048x2048xbf16> loc(#loc2555)
      %31885 = stablehlo.multiply %31884, %31884 : tensor<1x1x2048x2048xbf16> loc(#loc2556)
      %31886 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %31887 = stablehlo.broadcast_in_dim %31886, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2556)
      %31888 = stablehlo.multiply %31887, %31884 : tensor<1x1x2048x2048xbf16> loc(#loc2556)
      %31889 = stablehlo.convert %31885 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2550)
      %31890 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2551)
      %31891 = stablehlo.reduce(%31889 init: %31890) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2551)
        stablehlo.return %39096 : tensor<f32> loc(#loc2551)
      } loc(#loc2551)
      %31892 = stablehlo.broadcast_in_dim %31891, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2552)
      %31893 = stablehlo.convert %31892 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2553)
      %31894 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %31895 = stablehlo.broadcast_in_dim %31894, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2554)
      %31896 = stablehlo.divide %31893, %31895 : tensor<1x1x2048x1xbf16> loc(#loc2554)
      %31897 = stablehlo.broadcast_in_dim %31882, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2555)
      %31898 = stablehlo.subtract %31873, %31897 : tensor<1x1x2048x2048xbf16> loc(#loc2555)
      %31899 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %31900 = stablehlo.broadcast_in_dim %31899, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2557)
      %31901 = stablehlo.add %31896, %31900 : tensor<1x1x2048x1xbf16> loc(#loc2557)
      %31902 = stablehlo.rsqrt %31901 : tensor<1x1x2048x1xbf16> loc(#loc2558)
      %31903 = stablehlo.divide %31902, %31901 : tensor<1x1x2048x1xbf16> loc(#loc2559)
      %31904 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31905 = stablehlo.broadcast_in_dim %31904, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2560)
      %31906 = stablehlo.multiply %31905, %31903 : tensor<1x1x2048x1xbf16> loc(#loc2560)
      %31907 = stablehlo.broadcast_in_dim %31902, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2560)
      %31908 = stablehlo.multiply %31898, %31907 : tensor<1x1x2048x2048xbf16> loc(#loc2560)
      %31909 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31910 = stablehlo.broadcast_in_dim %31909, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2561)
      %31911 = stablehlo.add %31910, %30614#0 : tensor<1x2048xbf16> loc(#loc2561)
      %31912 = stablehlo.broadcast_in_dim %31911, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2562)
      %31913 = stablehlo.broadcast_in_dim %31912, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2563)
      %31914 = stablehlo.multiply %31908, %31913 : tensor<1x1x2048x2048xbf16> loc(#loc2563)
      %31915 = stablehlo.broadcast_in_dim %30613#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2564)
      %31916 = stablehlo.broadcast_in_dim %31915, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2565)
      %31917 = stablehlo.add %31914, %31916 : tensor<1x1x2048x2048xbf16> loc(#loc2565)
      %31918 = func.call @_einsum_683(%31917, %30610#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2566)
      %31919 = arith.select %ptrue, %31918, %31918 : tensor<1x1x2048x8192xbf16>
      %31920 = stablehlo.broadcast_in_dim %30609#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2568)
      %31921 = stablehlo.broadcast_in_dim %31920, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2569)
      %31922 = stablehlo.add %31919, %31921 : tensor<1x1x2048x8192xbf16> loc(#loc2569)
      %31923 = stablehlo.multiply %31922, %31922 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31924 = stablehlo.multiply %31922, %31923 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31925 = stablehlo.multiply %31922, %31922 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31926 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %31927 = stablehlo.broadcast_in_dim %31926, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31928 = stablehlo.multiply %31927, %31925 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31929 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %31930 = stablehlo.broadcast_in_dim %31929, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31931 = stablehlo.multiply %31930, %31924 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31932 = stablehlo.add %31922, %31931 : tensor<1x1x2048x8192xbf16> loc(#loc2571)
      %31933 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %31934 = stablehlo.broadcast_in_dim %31933, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31935 = stablehlo.multiply %31934, %31932 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31936 = stablehlo.tanh %31935 : tensor<1x1x2048x8192xbf16> loc(#loc2572)
      %31937 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31938 = stablehlo.broadcast_in_dim %31937, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2573)
      %31939 = stablehlo.subtract %31938, %31936 : tensor<1x1x2048x8192xbf16> loc(#loc2573)
      %31940 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31941 = stablehlo.broadcast_in_dim %31940, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2571)
      %31942 = stablehlo.add %31941, %31936 : tensor<1x1x2048x8192xbf16> loc(#loc2571)
      %31943 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %31944 = stablehlo.broadcast_in_dim %31943, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31945 = stablehlo.multiply %31944, %31942 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31946 = stablehlo.multiply %31922, %31945 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %31947 = stablehlo.convert %31874 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2574)
      %31948 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31949 = stablehlo.broadcast_in_dim %31948, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2575)
      %31950 = stablehlo.subtract %31949, %31947 : tensor<1x1x2048x1xf32> loc(#loc2575)
      %31951 = stablehlo.convert %31950 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2576)
      %31952 = stablehlo.broadcast_in_dim %31951, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2577)
      %31953 = stablehlo.multiply %31946, %31952 : tensor<1x1x2048x8192xbf16> loc(#loc2577)
      %31954 = func.call @_einsum_684(%31953, %30612#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2578)
      %31955 = arith.select %ptrue, %31954, %31954 : tensor<1x1x2048x2048xbf16>
      %31956 = stablehlo.broadcast_in_dim %30611#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2580)
      %31957 = stablehlo.broadcast_in_dim %31956, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2581)
      %31958 = stablehlo.add %31955, %31957 : tensor<1x1x2048x2048xbf16> loc(#loc2581)
      %31959 = stablehlo.convert %31874 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2574)
      %31960 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31961 = stablehlo.broadcast_in_dim %31960, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2575)
      %31962 = stablehlo.subtract %31961, %31959 : tensor<1x1x2048x1xf32> loc(#loc2575)
      %31963 = stablehlo.convert %31962 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2582)
      %31964 = stablehlo.broadcast_in_dim %31963, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2583)
      %31965 = stablehlo.multiply %31958, %31964 : tensor<1x1x2048x2048xbf16> loc(#loc2583)
      %31966 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %31967 = stablehlo.broadcast_in_dim %31966, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2584)
      %31968 = stablehlo.multiply %31965, %31967 : tensor<1x1x2048x2048xbf16> loc(#loc2584)
      %31969 = stablehlo.add %31873, %31968 : tensor<1x1x2048x2048xbf16> loc(#loc2585)
      %31970 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2586)
      %31971 = stablehlo.convert %31969 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2587)
      %31972 = stablehlo.convert %31970 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2588)
      %31973 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31974 = stablehlo.broadcast_in_dim %31973, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2589)
      %31975 = stablehlo.subtract %31974, %31972 : tensor<1x1x2048x1xf32> loc(#loc2589)
      %31976 = stablehlo.convert %31975 : tensor<1x1x2048x1xf32> loc(#loc2590)
      %31977 = stablehlo.broadcast_in_dim %31976, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %31978 = stablehlo.multiply %31971, %31977 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %31979 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2592)
      %31980 = stablehlo.reduce(%31978 init: %31979) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2592)
        stablehlo.return %39096 : tensor<f32> loc(#loc2592)
      } loc(#loc2592)
      %31981 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31982 = stablehlo.broadcast_in_dim %31981, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2593)
      %31983 = stablehlo.convert %31975 : tensor<1x1x2048x1xf32> loc(#loc2590)
      %31984 = stablehlo.broadcast_in_dim %31982, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2594)
      %31985 = stablehlo.broadcast_in_dim %31983, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %31986 = stablehlo.multiply %31984, %31985 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %31987 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2592)
      %31988 = stablehlo.reduce(%31986 init: %31987) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2592)
        stablehlo.return %39096 : tensor<f32> loc(#loc2592)
      } loc(#loc2592)
      %31989 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %31990 = stablehlo.broadcast_in_dim %31989, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2595)
      %31991 = stablehlo.maximum %31990, %31988 : tensor<1xf32> loc(#loc2595)
      %31992 = stablehlo.divide %31980, %31991 : tensor<1xf32> loc(#loc2596)
      %31993 = stablehlo.broadcast_in_dim %31992, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2597)
      %31994 = stablehlo.broadcast_in_dim %31993, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2598)
      %31995 = stablehlo.subtract %31971, %31994 : tensor<1x1x2048x2048xf32> loc(#loc2598)
      %31996 = stablehlo.convert %31975 : tensor<1x1x2048x1xf32> loc(#loc2590)
      %31997 = stablehlo.broadcast_in_dim %31996, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %31998 = stablehlo.multiply %31995, %31997 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %31999 = stablehlo.multiply %31998, %31998 : tensor<1x1x2048x2048xf32> loc(#loc2599)
      %32000 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %32001 = stablehlo.broadcast_in_dim %32000, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2599)
      %32002 = stablehlo.multiply %32001, %31998 : tensor<1x1x2048x2048xf32> loc(#loc2599)
      %32003 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2592)
      %32004 = stablehlo.reduce(%31999 init: %32003) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2592)
        stablehlo.return %39096 : tensor<f32> loc(#loc2592)
      } loc(#loc2592)
      %32005 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32006 = stablehlo.broadcast_in_dim %32005, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2595)
      %32007 = stablehlo.maximum %32006, %31988 : tensor<1xf32> loc(#loc2595)
      %32008 = stablehlo.divide %32004, %32007 : tensor<1xf32> loc(#loc2596)
      %32009 = stablehlo.sqrt %32008 : tensor<1xf32> loc(#loc2600)
      %32010 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32011 = stablehlo.broadcast_in_dim %32010, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2601)
      %32012 = stablehlo.divide %32011, %32009 : tensor<1xf32> loc(#loc2601)
      %32013 = stablehlo.convert %31975 : tensor<1x1x2048x1xf32> loc(#loc2590)
      %32014 = stablehlo.broadcast_in_dim %32013, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %32015 = stablehlo.multiply %31971, %32014 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %32016 = stablehlo.abs %32015 : tensor<1x1x2048x2048xf32> loc(#loc2602)
      %32017 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %32018 = stablehlo.broadcast_in_dim %32017, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2603)
      %32019 = stablehlo.compare  GE, %32015, %32018,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2603)
      %32020 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2604)
      %32021 = stablehlo.reduce(%32016 init: %32020) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2604)
        stablehlo.return %39096 : tensor<f32> loc(#loc2604)
      } loc(#loc2604)
      %32022 = stablehlo.reshape %32021 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2605)
      %32023 = stablehlo.broadcast_in_dim %32022, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2606)
      %32024 = stablehlo.compare  EQ, %32016, %32023,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2606)
      %32025 = stablehlo.convert %32024 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2607)
      %32026 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2608)
      %32027 = stablehlo.reduce(%32025 init: %32026) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2608)
        stablehlo.return %39096 : tensor<f32> loc(#loc2608)
      } loc(#loc2608)
      %32028 = stablehlo.convert %31969 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2609)
      %32029 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2610)
      %32030 = stablehlo.reduce(%32028 init: %32029) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2610)
        stablehlo.return %39096 : tensor<f32> loc(#loc2610)
      } loc(#loc2610)
      %32031 = stablehlo.broadcast_in_dim %32030, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2611)
      %32032 = stablehlo.convert %32031 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2612)
      %32033 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32034 = stablehlo.broadcast_in_dim %32033, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2613)
      %32035 = stablehlo.divide %32032, %32034 : tensor<1x1x2048x1xbf16> loc(#loc2613)
      %32036 = stablehlo.broadcast_in_dim %32035, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2614)
      %32037 = stablehlo.subtract %31969, %32036 : tensor<1x1x2048x2048xbf16> loc(#loc2614)
      %32038 = stablehlo.multiply %32037, %32037 : tensor<1x1x2048x2048xbf16> loc(#loc2615)
      %32039 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32040 = stablehlo.broadcast_in_dim %32039, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2615)
      %32041 = stablehlo.multiply %32040, %32037 : tensor<1x1x2048x2048xbf16> loc(#loc2615)
      %32042 = stablehlo.convert %32038 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2609)
      %32043 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2610)
      %32044 = stablehlo.reduce(%32042 init: %32043) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2610)
        stablehlo.return %39096 : tensor<f32> loc(#loc2610)
      } loc(#loc2610)
      %32045 = stablehlo.broadcast_in_dim %32044, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2611)
      %32046 = stablehlo.convert %32045 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2612)
      %32047 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32048 = stablehlo.broadcast_in_dim %32047, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2613)
      %32049 = stablehlo.divide %32046, %32048 : tensor<1x1x2048x1xbf16> loc(#loc2613)
      %32050 = stablehlo.broadcast_in_dim %32035, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2614)
      %32051 = stablehlo.subtract %31969, %32050 : tensor<1x1x2048x2048xbf16> loc(#loc2614)
      %32052 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32053 = stablehlo.broadcast_in_dim %32052, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2616)
      %32054 = stablehlo.add %32049, %32053 : tensor<1x1x2048x1xbf16> loc(#loc2616)
      %32055 = stablehlo.rsqrt %32054 : tensor<1x1x2048x1xbf16> loc(#loc2617)
      %32056 = stablehlo.divide %32055, %32054 : tensor<1x1x2048x1xbf16> loc(#loc2618)
      %32057 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32058 = stablehlo.broadcast_in_dim %32057, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2619)
      %32059 = stablehlo.multiply %32058, %32056 : tensor<1x1x2048x1xbf16> loc(#loc2619)
      %32060 = stablehlo.broadcast_in_dim %32055, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2619)
      %32061 = stablehlo.multiply %32051, %32060 : tensor<1x1x2048x2048xbf16> loc(#loc2619)
      %32062 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32063 = stablehlo.broadcast_in_dim %32062, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2620)
      %32064 = stablehlo.add %32063, %30627#0 : tensor<1x2048xbf16> loc(#loc2620)
      %32065 = stablehlo.broadcast_in_dim %32064, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2621)
      %32066 = stablehlo.broadcast_in_dim %32065, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2622)
      %32067 = stablehlo.multiply %32061, %32066 : tensor<1x1x2048x2048xbf16> loc(#loc2622)
      %32068 = stablehlo.broadcast_in_dim %30626#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2623)
      %32069 = stablehlo.broadcast_in_dim %32068, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2624)
      %32070 = stablehlo.add %32067, %32069 : tensor<1x1x2048x2048xbf16> loc(#loc2624)
      %32071 = func.call @_einsum_685(%32070, %30628#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2625)
      %32072 = "stablehlo.slice"(%32071) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2626)
      %32073 = stablehlo.reshape %32072 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2627)
      %32074 = "stablehlo.slice"(%32071) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2628)
      %32075 = stablehlo.reshape %32074 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2627)
      %32076 = "stablehlo.slice"(%32071) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2629)
      %32077 = stablehlo.reshape %32076 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2627)
      %32078 = arith.select %ptrue, %32073, %32073 : tensor<1x1x2048x32x64xbf16>
      %32079 = arith.select %ptrue, %32075, %32075 : tensor<1x1x2048x32x64xbf16>
      %32080 = arith.select %ptrue, %32077, %32077 : tensor<1x1x2048x32x64xbf16>
      %32081:2 = func.call @softplus_686(%30629#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2631)
      %32082 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %32083 = stablehlo.broadcast_in_dim %32082, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2632)
      %32084 = stablehlo.multiply %32083, %32081#0 : tensor<1x64xbf16> loc(#loc2632)
      %32085 = stablehlo.broadcast_in_dim %32084, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2633)
      %32086 = stablehlo.broadcast_in_dim %32085, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2634)
      %32087 = stablehlo.multiply %32078, %32086 : tensor<1x1x2048x32x64xbf16> loc(#loc2634)
      %32088 = func.call @_einsum_688(%32087, %32079) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2635)
      %32089 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32090 = stablehlo.broadcast_in_dim %32089, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2636)
      %32091 = stablehlo.divide %32088, %32090 : tensor<1x1x32x2048x2048xbf16> loc(#loc2636)
      %32092 = stablehlo.tanh %32091 : tensor<1x1x32x2048x2048xbf16> loc(#loc2637)
      %32093 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32094 = stablehlo.broadcast_in_dim %32093, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2638)
      %32095 = stablehlo.subtract %32094, %32092 : tensor<1x1x32x2048x2048xbf16> loc(#loc2638)
      %32096 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32097 = stablehlo.broadcast_in_dim %32096, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2639)
      %32098 = stablehlo.multiply %32097, %32092 : tensor<1x1x32x2048x2048xbf16> loc(#loc2639)
      %32099 = stablehlo.convert %32098 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2640)
      %32100 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32101 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32102 = stablehlo.multiply %32100, %32101 : tensor<f32> loc(#loc2641)
      %32103 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2642)
      %32104 = stablehlo.broadcast_in_dim %32102, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2643)
      %32105 = stablehlo.compare  GE, %32103, %32104,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2643)
      %32106 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32107:2 = func.call @_where_689(%32105, %32099, %32106) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2644)
      %32108 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2645)
      %32109 = stablehlo.reduce(%32107#0 init: %32108) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2645)
        stablehlo.return %39096 : tensor<f32> loc(#loc2645)
      } loc(#loc2645)
      %32110 = stablehlo.broadcast_in_dim %32109, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2646)
      %32111 = stablehlo.broadcast_in_dim %32110, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2647)
      %32112 = stablehlo.subtract %32107#0, %32111 : tensor<1x1x32x2048x2048xf32> loc(#loc2647)
      %32113 = stablehlo.exponential %32112 : tensor<1x1x32x2048x2048xf32> loc(#loc2648)
      %32114 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2649)
      %32115 = stablehlo.reduce(%32113 init: %32114) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2649)
        stablehlo.return %39096 : tensor<f32> loc(#loc2649)
      } loc(#loc2649)
      %32116 = stablehlo.broadcast_in_dim %32115, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2646)
      %32117 = stablehlo.broadcast_in_dim %32116, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2650)
      %32118 = stablehlo.divide %32113, %32117 : tensor<1x1x32x2048x2048xf32> loc(#loc2650)
      %32119 = stablehlo.multiply %32116, %32116 : tensor<1x1x32x2048x1xf32> loc(#loc2651)
      %32120 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2652)
      %32121 = stablehlo.broadcast_in_dim %32120, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2653)
      %32122 = stablehlo.divide %32121, %32119 : tensor<1x1x32x2048x1xf32> loc(#loc2650)
      %32123 = stablehlo.convert %32118 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2654)
      %32124 = func.call @_einsum_690(%32123, %32080) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2655)
      %32125 = arith.select %ptrue, %32124, %32124 : tensor<1x1x2048x32x64xbf16>
      %32126 = func.call @_einsum_691(%32125, %30630#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2656)
      %32127 = arith.select %ptrue, %32126, %32126 : tensor<1x1x2048x2048xbf16>
      %32128 = stablehlo.add %32127, %31969 : tensor<1x1x2048x2048xbf16> loc(#loc2658)
      %32129 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2659)
      %32130 = stablehlo.convert %32128 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2660)
      %32131 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2661)
      %32132 = stablehlo.reduce(%32130 init: %32131) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2661)
        stablehlo.return %39096 : tensor<f32> loc(#loc2661)
      } loc(#loc2661)
      %32133 = stablehlo.broadcast_in_dim %32132, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2662)
      %32134 = stablehlo.convert %32133 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2663)
      %32135 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32136 = stablehlo.broadcast_in_dim %32135, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2664)
      %32137 = stablehlo.divide %32134, %32136 : tensor<1x1x2048x1xbf16> loc(#loc2664)
      %32138 = stablehlo.broadcast_in_dim %32137, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2665)
      %32139 = stablehlo.subtract %32128, %32138 : tensor<1x1x2048x2048xbf16> loc(#loc2665)
      %32140 = stablehlo.multiply %32139, %32139 : tensor<1x1x2048x2048xbf16> loc(#loc2666)
      %32141 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32142 = stablehlo.broadcast_in_dim %32141, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2666)
      %32143 = stablehlo.multiply %32142, %32139 : tensor<1x1x2048x2048xbf16> loc(#loc2666)
      %32144 = stablehlo.convert %32140 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2660)
      %32145 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2661)
      %32146 = stablehlo.reduce(%32144 init: %32145) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2661)
        stablehlo.return %39096 : tensor<f32> loc(#loc2661)
      } loc(#loc2661)
      %32147 = stablehlo.broadcast_in_dim %32146, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2662)
      %32148 = stablehlo.convert %32147 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2663)
      %32149 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32150 = stablehlo.broadcast_in_dim %32149, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2664)
      %32151 = stablehlo.divide %32148, %32150 : tensor<1x1x2048x1xbf16> loc(#loc2664)
      %32152 = stablehlo.broadcast_in_dim %32137, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2665)
      %32153 = stablehlo.subtract %32128, %32152 : tensor<1x1x2048x2048xbf16> loc(#loc2665)
      %32154 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32155 = stablehlo.broadcast_in_dim %32154, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2667)
      %32156 = stablehlo.add %32151, %32155 : tensor<1x1x2048x1xbf16> loc(#loc2667)
      %32157 = stablehlo.rsqrt %32156 : tensor<1x1x2048x1xbf16> loc(#loc2668)
      %32158 = stablehlo.divide %32157, %32156 : tensor<1x1x2048x1xbf16> loc(#loc2669)
      %32159 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32160 = stablehlo.broadcast_in_dim %32159, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2670)
      %32161 = stablehlo.multiply %32160, %32158 : tensor<1x1x2048x1xbf16> loc(#loc2670)
      %32162 = stablehlo.broadcast_in_dim %32157, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2670)
      %32163 = stablehlo.multiply %32153, %32162 : tensor<1x1x2048x2048xbf16> loc(#loc2670)
      %32164 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32165 = stablehlo.broadcast_in_dim %32164, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2671)
      %32166 = stablehlo.add %32165, %30625#0 : tensor<1x2048xbf16> loc(#loc2671)
      %32167 = stablehlo.broadcast_in_dim %32166, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2672)
      %32168 = stablehlo.broadcast_in_dim %32167, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2673)
      %32169 = stablehlo.multiply %32163, %32168 : tensor<1x1x2048x2048xbf16> loc(#loc2673)
      %32170 = stablehlo.broadcast_in_dim %30624#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2674)
      %32171 = stablehlo.broadcast_in_dim %32170, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2675)
      %32172 = stablehlo.add %32169, %32171 : tensor<1x1x2048x2048xbf16> loc(#loc2675)
      %32173 = func.call @_einsum_692(%32172, %30621#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2676)
      %32174 = arith.select %ptrue, %32173, %32173 : tensor<1x1x2048x8192xbf16>
      %32175 = stablehlo.broadcast_in_dim %30620#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2678)
      %32176 = stablehlo.broadcast_in_dim %32175, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2679)
      %32177 = stablehlo.add %32174, %32176 : tensor<1x1x2048x8192xbf16> loc(#loc2679)
      %32178 = stablehlo.multiply %32177, %32177 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32179 = stablehlo.multiply %32177, %32178 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32180 = stablehlo.multiply %32177, %32177 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32181 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %32182 = stablehlo.broadcast_in_dim %32181, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32183 = stablehlo.multiply %32182, %32180 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32184 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %32185 = stablehlo.broadcast_in_dim %32184, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32186 = stablehlo.multiply %32185, %32179 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32187 = stablehlo.add %32177, %32186 : tensor<1x1x2048x8192xbf16> loc(#loc2681)
      %32188 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %32189 = stablehlo.broadcast_in_dim %32188, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32190 = stablehlo.multiply %32189, %32187 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32191 = stablehlo.tanh %32190 : tensor<1x1x2048x8192xbf16> loc(#loc2682)
      %32192 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32193 = stablehlo.broadcast_in_dim %32192, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2683)
      %32194 = stablehlo.subtract %32193, %32191 : tensor<1x1x2048x8192xbf16> loc(#loc2683)
      %32195 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32196 = stablehlo.broadcast_in_dim %32195, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2681)
      %32197 = stablehlo.add %32196, %32191 : tensor<1x1x2048x8192xbf16> loc(#loc2681)
      %32198 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32199 = stablehlo.broadcast_in_dim %32198, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32200 = stablehlo.multiply %32199, %32197 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32201 = stablehlo.multiply %32177, %32200 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %32202 = stablehlo.convert %32129 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2684)
      %32203 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32204 = stablehlo.broadcast_in_dim %32203, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2685)
      %32205 = stablehlo.subtract %32204, %32202 : tensor<1x1x2048x1xf32> loc(#loc2685)
      %32206 = stablehlo.convert %32205 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2686)
      %32207 = stablehlo.broadcast_in_dim %32206, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2687)
      %32208 = stablehlo.multiply %32201, %32207 : tensor<1x1x2048x8192xbf16> loc(#loc2687)
      %32209 = func.call @_einsum_693(%32208, %30623#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2688)
      %32210 = arith.select %ptrue, %32209, %32209 : tensor<1x1x2048x2048xbf16>
      %32211 = stablehlo.broadcast_in_dim %30622#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2690)
      %32212 = stablehlo.broadcast_in_dim %32211, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2691)
      %32213 = stablehlo.add %32210, %32212 : tensor<1x1x2048x2048xbf16> loc(#loc2691)
      %32214 = stablehlo.convert %32129 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2684)
      %32215 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32216 = stablehlo.broadcast_in_dim %32215, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2685)
      %32217 = stablehlo.subtract %32216, %32214 : tensor<1x1x2048x1xf32> loc(#loc2685)
      %32218 = stablehlo.convert %32217 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2692)
      %32219 = stablehlo.broadcast_in_dim %32218, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2693)
      %32220 = stablehlo.multiply %32213, %32219 : tensor<1x1x2048x2048xbf16> loc(#loc2693)
      %32221 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32222 = stablehlo.broadcast_in_dim %32221, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2694)
      %32223 = stablehlo.multiply %32220, %32222 : tensor<1x1x2048x2048xbf16> loc(#loc2694)
      %32224 = stablehlo.add %32128, %32223 : tensor<1x1x2048x2048xbf16> loc(#loc2695)
      %32225 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2696)
      %32226 = stablehlo.convert %32224 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2697)
      %32227 = stablehlo.convert %32225 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2698)
      %32228 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32229 = stablehlo.broadcast_in_dim %32228, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2699)
      %32230 = stablehlo.subtract %32229, %32227 : tensor<1x1x2048x1xf32> loc(#loc2699)
      %32231 = stablehlo.convert %32230 : tensor<1x1x2048x1xf32> loc(#loc2700)
      %32232 = stablehlo.broadcast_in_dim %32231, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32233 = stablehlo.multiply %32226, %32232 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32234 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2702)
      %32235 = stablehlo.reduce(%32233 init: %32234) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2702)
        stablehlo.return %39096 : tensor<f32> loc(#loc2702)
      } loc(#loc2702)
      %32236 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32237 = stablehlo.broadcast_in_dim %32236, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2703)
      %32238 = stablehlo.convert %32230 : tensor<1x1x2048x1xf32> loc(#loc2700)
      %32239 = stablehlo.broadcast_in_dim %32237, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2704)
      %32240 = stablehlo.broadcast_in_dim %32238, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32241 = stablehlo.multiply %32239, %32240 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32242 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2702)
      %32243 = stablehlo.reduce(%32241 init: %32242) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2702)
        stablehlo.return %39096 : tensor<f32> loc(#loc2702)
      } loc(#loc2702)
      %32244 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32245 = stablehlo.broadcast_in_dim %32244, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2705)
      %32246 = stablehlo.maximum %32245, %32243 : tensor<1xf32> loc(#loc2705)
      %32247 = stablehlo.divide %32235, %32246 : tensor<1xf32> loc(#loc2706)
      %32248 = stablehlo.broadcast_in_dim %32247, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2707)
      %32249 = stablehlo.broadcast_in_dim %32248, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2708)
      %32250 = stablehlo.subtract %32226, %32249 : tensor<1x1x2048x2048xf32> loc(#loc2708)
      %32251 = stablehlo.convert %32230 : tensor<1x1x2048x1xf32> loc(#loc2700)
      %32252 = stablehlo.broadcast_in_dim %32251, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32253 = stablehlo.multiply %32250, %32252 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32254 = stablehlo.multiply %32253, %32253 : tensor<1x1x2048x2048xf32> loc(#loc2709)
      %32255 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %32256 = stablehlo.broadcast_in_dim %32255, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2709)
      %32257 = stablehlo.multiply %32256, %32253 : tensor<1x1x2048x2048xf32> loc(#loc2709)
      %32258 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2702)
      %32259 = stablehlo.reduce(%32254 init: %32258) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2702)
        stablehlo.return %39096 : tensor<f32> loc(#loc2702)
      } loc(#loc2702)
      %32260 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32261 = stablehlo.broadcast_in_dim %32260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2705)
      %32262 = stablehlo.maximum %32261, %32243 : tensor<1xf32> loc(#loc2705)
      %32263 = stablehlo.divide %32259, %32262 : tensor<1xf32> loc(#loc2706)
      %32264 = stablehlo.sqrt %32263 : tensor<1xf32> loc(#loc2710)
      %32265 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32266 = stablehlo.broadcast_in_dim %32265, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2711)
      %32267 = stablehlo.divide %32266, %32264 : tensor<1xf32> loc(#loc2711)
      %32268 = stablehlo.convert %32230 : tensor<1x1x2048x1xf32> loc(#loc2700)
      %32269 = stablehlo.broadcast_in_dim %32268, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32270 = stablehlo.multiply %32226, %32269 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %32271 = stablehlo.abs %32270 : tensor<1x1x2048x2048xf32> loc(#loc2712)
      %32272 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %32273 = stablehlo.broadcast_in_dim %32272, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2713)
      %32274 = stablehlo.compare  GE, %32270, %32273,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2713)
      %32275 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2714)
      %32276 = stablehlo.reduce(%32271 init: %32275) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2714)
        stablehlo.return %39096 : tensor<f32> loc(#loc2714)
      } loc(#loc2714)
      %32277 = stablehlo.reshape %32276 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2715)
      %32278 = stablehlo.broadcast_in_dim %32277, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2716)
      %32279 = stablehlo.compare  EQ, %32271, %32278,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2716)
      %32280 = stablehlo.convert %32279 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2717)
      %32281 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2718)
      %32282 = stablehlo.reduce(%32280 init: %32281) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2718)
        stablehlo.return %39096 : tensor<f32> loc(#loc2718)
      } loc(#loc2718)
      %32283 = stablehlo.convert %32224 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2719)
      %32284 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2720)
      %32285 = stablehlo.reduce(%32283 init: %32284) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2720)
        stablehlo.return %39096 : tensor<f32> loc(#loc2720)
      } loc(#loc2720)
      %32286 = stablehlo.broadcast_in_dim %32285, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2721)
      %32287 = stablehlo.convert %32286 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2722)
      %32288 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32289 = stablehlo.broadcast_in_dim %32288, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2723)
      %32290 = stablehlo.divide %32287, %32289 : tensor<1x1x2048x1xbf16> loc(#loc2723)
      %32291 = stablehlo.broadcast_in_dim %32290, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2724)
      %32292 = stablehlo.subtract %32224, %32291 : tensor<1x1x2048x2048xbf16> loc(#loc2724)
      %32293 = stablehlo.multiply %32292, %32292 : tensor<1x1x2048x2048xbf16> loc(#loc2725)
      %32294 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32295 = stablehlo.broadcast_in_dim %32294, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2725)
      %32296 = stablehlo.multiply %32295, %32292 : tensor<1x1x2048x2048xbf16> loc(#loc2725)
      %32297 = stablehlo.convert %32293 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2719)
      %32298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2720)
      %32299 = stablehlo.reduce(%32297 init: %32298) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2720)
        stablehlo.return %39096 : tensor<f32> loc(#loc2720)
      } loc(#loc2720)
      %32300 = stablehlo.broadcast_in_dim %32299, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2721)
      %32301 = stablehlo.convert %32300 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2722)
      %32302 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32303 = stablehlo.broadcast_in_dim %32302, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2723)
      %32304 = stablehlo.divide %32301, %32303 : tensor<1x1x2048x1xbf16> loc(#loc2723)
      %32305 = stablehlo.broadcast_in_dim %32290, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2724)
      %32306 = stablehlo.subtract %32224, %32305 : tensor<1x1x2048x2048xbf16> loc(#loc2724)
      %32307 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32308 = stablehlo.broadcast_in_dim %32307, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2726)
      %32309 = stablehlo.add %32304, %32308 : tensor<1x1x2048x1xbf16> loc(#loc2726)
      %32310 = stablehlo.rsqrt %32309 : tensor<1x1x2048x1xbf16> loc(#loc2727)
      %32311 = stablehlo.divide %32310, %32309 : tensor<1x1x2048x1xbf16> loc(#loc2728)
      %32312 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32313 = stablehlo.broadcast_in_dim %32312, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2729)
      %32314 = stablehlo.multiply %32313, %32311 : tensor<1x1x2048x1xbf16> loc(#loc2729)
      %32315 = stablehlo.broadcast_in_dim %32310, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2729)
      %32316 = stablehlo.multiply %32306, %32315 : tensor<1x1x2048x2048xbf16> loc(#loc2729)
      %32317 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32318 = stablehlo.broadcast_in_dim %32317, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2730)
      %32319 = stablehlo.add %32318, %30638#0 : tensor<1x2048xbf16> loc(#loc2730)
      %32320 = stablehlo.broadcast_in_dim %32319, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2731)
      %32321 = stablehlo.broadcast_in_dim %32320, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2732)
      %32322 = stablehlo.multiply %32316, %32321 : tensor<1x1x2048x2048xbf16> loc(#loc2732)
      %32323 = stablehlo.broadcast_in_dim %30637#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2733)
      %32324 = stablehlo.broadcast_in_dim %32323, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2734)
      %32325 = stablehlo.add %32322, %32324 : tensor<1x1x2048x2048xbf16> loc(#loc2734)
      %32326 = func.call @_einsum_694(%32325, %30639#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2735)
      %32327 = "stablehlo.slice"(%32326) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2736)
      %32328 = stablehlo.reshape %32327 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2737)
      %32329 = "stablehlo.slice"(%32326) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2738)
      %32330 = stablehlo.reshape %32329 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2737)
      %32331 = "stablehlo.slice"(%32326) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2739)
      %32332 = stablehlo.reshape %32331 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2737)
      %32333 = arith.select %ptrue, %32328, %32328 : tensor<1x1x2048x32x64xbf16>
      %32334 = arith.select %ptrue, %32330, %32330 : tensor<1x1x2048x32x64xbf16>
      %32335 = arith.select %ptrue, %32332, %32332 : tensor<1x1x2048x32x64xbf16>
      %32336:2 = func.call @softplus_695(%30640#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2741)
      %32337 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %32338 = stablehlo.broadcast_in_dim %32337, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2742)
      %32339 = stablehlo.multiply %32338, %32336#0 : tensor<1x64xbf16> loc(#loc2742)
      %32340 = stablehlo.broadcast_in_dim %32339, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2743)
      %32341 = stablehlo.broadcast_in_dim %32340, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2744)
      %32342 = stablehlo.multiply %32333, %32341 : tensor<1x1x2048x32x64xbf16> loc(#loc2744)
      %32343 = func.call @_einsum_697(%32342, %32334) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2745)
      %32344 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32345 = stablehlo.broadcast_in_dim %32344, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2746)
      %32346 = stablehlo.divide %32343, %32345 : tensor<1x1x32x2048x2048xbf16> loc(#loc2746)
      %32347 = stablehlo.tanh %32346 : tensor<1x1x32x2048x2048xbf16> loc(#loc2747)
      %32348 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32349 = stablehlo.broadcast_in_dim %32348, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2748)
      %32350 = stablehlo.subtract %32349, %32347 : tensor<1x1x32x2048x2048xbf16> loc(#loc2748)
      %32351 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32352 = stablehlo.broadcast_in_dim %32351, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2749)
      %32353 = stablehlo.multiply %32352, %32347 : tensor<1x1x32x2048x2048xbf16> loc(#loc2749)
      %32354 = stablehlo.convert %32353 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2750)
      %32355 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32356 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32357 = stablehlo.multiply %32355, %32356 : tensor<f32> loc(#loc2751)
      %32358 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2752)
      %32359 = stablehlo.broadcast_in_dim %32357, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2753)
      %32360 = stablehlo.compare  GE, %32358, %32359,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2753)
      %32361 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32362:2 = func.call @_where_698(%32360, %32354, %32361) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2754)
      %32363 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2755)
      %32364 = stablehlo.reduce(%32362#0 init: %32363) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2755)
        stablehlo.return %39096 : tensor<f32> loc(#loc2755)
      } loc(#loc2755)
      %32365 = stablehlo.broadcast_in_dim %32364, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2756)
      %32366 = stablehlo.broadcast_in_dim %32365, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2757)
      %32367 = stablehlo.subtract %32362#0, %32366 : tensor<1x1x32x2048x2048xf32> loc(#loc2757)
      %32368 = stablehlo.exponential %32367 : tensor<1x1x32x2048x2048xf32> loc(#loc2758)
      %32369 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2759)
      %32370 = stablehlo.reduce(%32368 init: %32369) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2759)
        stablehlo.return %39096 : tensor<f32> loc(#loc2759)
      } loc(#loc2759)
      %32371 = stablehlo.broadcast_in_dim %32370, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2756)
      %32372 = stablehlo.broadcast_in_dim %32371, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2760)
      %32373 = stablehlo.divide %32368, %32372 : tensor<1x1x32x2048x2048xf32> loc(#loc2760)
      %32374 = stablehlo.multiply %32371, %32371 : tensor<1x1x32x2048x1xf32> loc(#loc2761)
      %32375 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2762)
      %32376 = stablehlo.broadcast_in_dim %32375, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2763)
      %32377 = stablehlo.divide %32376, %32374 : tensor<1x1x32x2048x1xf32> loc(#loc2760)
      %32378 = stablehlo.convert %32373 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2764)
      %32379 = func.call @_einsum_699(%32378, %32335) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2765)
      %32380 = arith.select %ptrue, %32379, %32379 : tensor<1x1x2048x32x64xbf16>
      %32381 = func.call @_einsum_700(%32380, %30641#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2766)
      %32382 = arith.select %ptrue, %32381, %32381 : tensor<1x1x2048x2048xbf16>
      %32383 = stablehlo.add %32382, %32224 : tensor<1x1x2048x2048xbf16> loc(#loc2768)
      %32384 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2769)
      %32385 = stablehlo.convert %32383 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2770)
      %32386 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2771)
      %32387 = stablehlo.reduce(%32385 init: %32386) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2771)
        stablehlo.return %39096 : tensor<f32> loc(#loc2771)
      } loc(#loc2771)
      %32388 = stablehlo.broadcast_in_dim %32387, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2772)
      %32389 = stablehlo.convert %32388 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2773)
      %32390 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32391 = stablehlo.broadcast_in_dim %32390, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2774)
      %32392 = stablehlo.divide %32389, %32391 : tensor<1x1x2048x1xbf16> loc(#loc2774)
      %32393 = stablehlo.broadcast_in_dim %32392, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2775)
      %32394 = stablehlo.subtract %32383, %32393 : tensor<1x1x2048x2048xbf16> loc(#loc2775)
      %32395 = stablehlo.multiply %32394, %32394 : tensor<1x1x2048x2048xbf16> loc(#loc2776)
      %32396 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32397 = stablehlo.broadcast_in_dim %32396, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2776)
      %32398 = stablehlo.multiply %32397, %32394 : tensor<1x1x2048x2048xbf16> loc(#loc2776)
      %32399 = stablehlo.convert %32395 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2770)
      %32400 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2771)
      %32401 = stablehlo.reduce(%32399 init: %32400) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2771)
        stablehlo.return %39096 : tensor<f32> loc(#loc2771)
      } loc(#loc2771)
      %32402 = stablehlo.broadcast_in_dim %32401, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2772)
      %32403 = stablehlo.convert %32402 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2773)
      %32404 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32405 = stablehlo.broadcast_in_dim %32404, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2774)
      %32406 = stablehlo.divide %32403, %32405 : tensor<1x1x2048x1xbf16> loc(#loc2774)
      %32407 = stablehlo.broadcast_in_dim %32392, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2775)
      %32408 = stablehlo.subtract %32383, %32407 : tensor<1x1x2048x2048xbf16> loc(#loc2775)
      %32409 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32410 = stablehlo.broadcast_in_dim %32409, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2777)
      %32411 = stablehlo.add %32406, %32410 : tensor<1x1x2048x1xbf16> loc(#loc2777)
      %32412 = stablehlo.rsqrt %32411 : tensor<1x1x2048x1xbf16> loc(#loc2778)
      %32413 = stablehlo.divide %32412, %32411 : tensor<1x1x2048x1xbf16> loc(#loc2779)
      %32414 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32415 = stablehlo.broadcast_in_dim %32414, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2780)
      %32416 = stablehlo.multiply %32415, %32413 : tensor<1x1x2048x1xbf16> loc(#loc2780)
      %32417 = stablehlo.broadcast_in_dim %32412, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2780)
      %32418 = stablehlo.multiply %32408, %32417 : tensor<1x1x2048x2048xbf16> loc(#loc2780)
      %32419 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32420 = stablehlo.broadcast_in_dim %32419, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2781)
      %32421 = stablehlo.add %32420, %30636#0 : tensor<1x2048xbf16> loc(#loc2781)
      %32422 = stablehlo.broadcast_in_dim %32421, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2782)
      %32423 = stablehlo.broadcast_in_dim %32422, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2783)
      %32424 = stablehlo.multiply %32418, %32423 : tensor<1x1x2048x2048xbf16> loc(#loc2783)
      %32425 = stablehlo.broadcast_in_dim %30635#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2784)
      %32426 = stablehlo.broadcast_in_dim %32425, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2785)
      %32427 = stablehlo.add %32424, %32426 : tensor<1x1x2048x2048xbf16> loc(#loc2785)
      %32428 = func.call @_einsum_701(%32427, %30632#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2786)
      %32429 = arith.select %ptrue, %32428, %32428 : tensor<1x1x2048x8192xbf16>
      %32430 = stablehlo.broadcast_in_dim %30631#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2788)
      %32431 = stablehlo.broadcast_in_dim %32430, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2789)
      %32432 = stablehlo.add %32429, %32431 : tensor<1x1x2048x8192xbf16> loc(#loc2789)
      %32433 = stablehlo.multiply %32432, %32432 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32434 = stablehlo.multiply %32432, %32433 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32435 = stablehlo.multiply %32432, %32432 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32436 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %32437 = stablehlo.broadcast_in_dim %32436, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32438 = stablehlo.multiply %32437, %32435 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32439 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %32440 = stablehlo.broadcast_in_dim %32439, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32441 = stablehlo.multiply %32440, %32434 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32442 = stablehlo.add %32432, %32441 : tensor<1x1x2048x8192xbf16> loc(#loc2791)
      %32443 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %32444 = stablehlo.broadcast_in_dim %32443, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32445 = stablehlo.multiply %32444, %32442 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32446 = stablehlo.tanh %32445 : tensor<1x1x2048x8192xbf16> loc(#loc2792)
      %32447 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32448 = stablehlo.broadcast_in_dim %32447, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2793)
      %32449 = stablehlo.subtract %32448, %32446 : tensor<1x1x2048x8192xbf16> loc(#loc2793)
      %32450 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32451 = stablehlo.broadcast_in_dim %32450, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2791)
      %32452 = stablehlo.add %32451, %32446 : tensor<1x1x2048x8192xbf16> loc(#loc2791)
      %32453 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32454 = stablehlo.broadcast_in_dim %32453, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32455 = stablehlo.multiply %32454, %32452 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32456 = stablehlo.multiply %32432, %32455 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %32457 = stablehlo.convert %32384 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2794)
      %32458 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32459 = stablehlo.broadcast_in_dim %32458, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2795)
      %32460 = stablehlo.subtract %32459, %32457 : tensor<1x1x2048x1xf32> loc(#loc2795)
      %32461 = stablehlo.convert %32460 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2796)
      %32462 = stablehlo.broadcast_in_dim %32461, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2797)
      %32463 = stablehlo.multiply %32456, %32462 : tensor<1x1x2048x8192xbf16> loc(#loc2797)
      %32464 = func.call @_einsum_702(%32463, %30634#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2798)
      %32465 = arith.select %ptrue, %32464, %32464 : tensor<1x1x2048x2048xbf16>
      %32466 = stablehlo.broadcast_in_dim %30633#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2800)
      %32467 = stablehlo.broadcast_in_dim %32466, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2801)
      %32468 = stablehlo.add %32465, %32467 : tensor<1x1x2048x2048xbf16> loc(#loc2801)
      %32469 = stablehlo.convert %32384 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2794)
      %32470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32471 = stablehlo.broadcast_in_dim %32470, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2795)
      %32472 = stablehlo.subtract %32471, %32469 : tensor<1x1x2048x1xf32> loc(#loc2795)
      %32473 = stablehlo.convert %32472 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2802)
      %32474 = stablehlo.broadcast_in_dim %32473, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2803)
      %32475 = stablehlo.multiply %32468, %32474 : tensor<1x1x2048x2048xbf16> loc(#loc2803)
      %32476 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32477 = stablehlo.broadcast_in_dim %32476, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2804)
      %32478 = stablehlo.multiply %32475, %32477 : tensor<1x1x2048x2048xbf16> loc(#loc2804)
      %32479 = stablehlo.add %32383, %32478 : tensor<1x1x2048x2048xbf16> loc(#loc2805)
      %32480 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2806)
      %32481 = stablehlo.convert %32479 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2807)
      %32482 = stablehlo.convert %32480 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2808)
      %32483 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32484 = stablehlo.broadcast_in_dim %32483, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2809)
      %32485 = stablehlo.subtract %32484, %32482 : tensor<1x1x2048x1xf32> loc(#loc2809)
      %32486 = stablehlo.convert %32485 : tensor<1x1x2048x1xf32> loc(#loc2810)
      %32487 = stablehlo.broadcast_in_dim %32486, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32488 = stablehlo.multiply %32481, %32487 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32489 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2812)
      %32490 = stablehlo.reduce(%32488 init: %32489) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2812)
        stablehlo.return %39096 : tensor<f32> loc(#loc2812)
      } loc(#loc2812)
      %32491 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32492 = stablehlo.broadcast_in_dim %32491, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2813)
      %32493 = stablehlo.convert %32485 : tensor<1x1x2048x1xf32> loc(#loc2810)
      %32494 = stablehlo.broadcast_in_dim %32492, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2814)
      %32495 = stablehlo.broadcast_in_dim %32493, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32496 = stablehlo.multiply %32494, %32495 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32497 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2812)
      %32498 = stablehlo.reduce(%32496 init: %32497) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2812)
        stablehlo.return %39096 : tensor<f32> loc(#loc2812)
      } loc(#loc2812)
      %32499 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32500 = stablehlo.broadcast_in_dim %32499, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2815)
      %32501 = stablehlo.maximum %32500, %32498 : tensor<1xf32> loc(#loc2815)
      %32502 = stablehlo.divide %32490, %32501 : tensor<1xf32> loc(#loc2816)
      %32503 = stablehlo.broadcast_in_dim %32502, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2817)
      %32504 = stablehlo.broadcast_in_dim %32503, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2818)
      %32505 = stablehlo.subtract %32481, %32504 : tensor<1x1x2048x2048xf32> loc(#loc2818)
      %32506 = stablehlo.convert %32485 : tensor<1x1x2048x1xf32> loc(#loc2810)
      %32507 = stablehlo.broadcast_in_dim %32506, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32508 = stablehlo.multiply %32505, %32507 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32509 = stablehlo.multiply %32508, %32508 : tensor<1x1x2048x2048xf32> loc(#loc2819)
      %32510 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %32511 = stablehlo.broadcast_in_dim %32510, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2819)
      %32512 = stablehlo.multiply %32511, %32508 : tensor<1x1x2048x2048xf32> loc(#loc2819)
      %32513 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2812)
      %32514 = stablehlo.reduce(%32509 init: %32513) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2812)
        stablehlo.return %39096 : tensor<f32> loc(#loc2812)
      } loc(#loc2812)
      %32515 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32516 = stablehlo.broadcast_in_dim %32515, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2815)
      %32517 = stablehlo.maximum %32516, %32498 : tensor<1xf32> loc(#loc2815)
      %32518 = stablehlo.divide %32514, %32517 : tensor<1xf32> loc(#loc2816)
      %32519 = stablehlo.sqrt %32518 : tensor<1xf32> loc(#loc2820)
      %32520 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32521 = stablehlo.broadcast_in_dim %32520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2821)
      %32522 = stablehlo.divide %32521, %32519 : tensor<1xf32> loc(#loc2821)
      %32523 = stablehlo.convert %32485 : tensor<1x1x2048x1xf32> loc(#loc2810)
      %32524 = stablehlo.broadcast_in_dim %32523, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32525 = stablehlo.multiply %32481, %32524 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %32526 = stablehlo.abs %32525 : tensor<1x1x2048x2048xf32> loc(#loc2822)
      %32527 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %32528 = stablehlo.broadcast_in_dim %32527, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2823)
      %32529 = stablehlo.compare  GE, %32525, %32528,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2823)
      %32530 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2824)
      %32531 = stablehlo.reduce(%32526 init: %32530) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2824)
        stablehlo.return %39096 : tensor<f32> loc(#loc2824)
      } loc(#loc2824)
      %32532 = stablehlo.reshape %32531 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2825)
      %32533 = stablehlo.broadcast_in_dim %32532, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2826)
      %32534 = stablehlo.compare  EQ, %32526, %32533,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2826)
      %32535 = stablehlo.convert %32534 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2827)
      %32536 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2828)
      %32537 = stablehlo.reduce(%32535 init: %32536) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2828)
        stablehlo.return %39096 : tensor<f32> loc(#loc2828)
      } loc(#loc2828)
      %32538 = stablehlo.convert %32479 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2829)
      %32539 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2830)
      %32540 = stablehlo.reduce(%32538 init: %32539) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2830)
        stablehlo.return %39096 : tensor<f32> loc(#loc2830)
      } loc(#loc2830)
      %32541 = stablehlo.broadcast_in_dim %32540, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2831)
      %32542 = stablehlo.convert %32541 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2832)
      %32543 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32544 = stablehlo.broadcast_in_dim %32543, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2833)
      %32545 = stablehlo.divide %32542, %32544 : tensor<1x1x2048x1xbf16> loc(#loc2833)
      %32546 = stablehlo.broadcast_in_dim %32545, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2834)
      %32547 = stablehlo.subtract %32479, %32546 : tensor<1x1x2048x2048xbf16> loc(#loc2834)
      %32548 = stablehlo.multiply %32547, %32547 : tensor<1x1x2048x2048xbf16> loc(#loc2835)
      %32549 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32550 = stablehlo.broadcast_in_dim %32549, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2835)
      %32551 = stablehlo.multiply %32550, %32547 : tensor<1x1x2048x2048xbf16> loc(#loc2835)
      %32552 = stablehlo.convert %32548 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2829)
      %32553 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2830)
      %32554 = stablehlo.reduce(%32552 init: %32553) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2830)
        stablehlo.return %39096 : tensor<f32> loc(#loc2830)
      } loc(#loc2830)
      %32555 = stablehlo.broadcast_in_dim %32554, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2831)
      %32556 = stablehlo.convert %32555 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2832)
      %32557 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32558 = stablehlo.broadcast_in_dim %32557, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2833)
      %32559 = stablehlo.divide %32556, %32558 : tensor<1x1x2048x1xbf16> loc(#loc2833)
      %32560 = stablehlo.broadcast_in_dim %32545, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2834)
      %32561 = stablehlo.subtract %32479, %32560 : tensor<1x1x2048x2048xbf16> loc(#loc2834)
      %32562 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32563 = stablehlo.broadcast_in_dim %32562, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2836)
      %32564 = stablehlo.add %32559, %32563 : tensor<1x1x2048x1xbf16> loc(#loc2836)
      %32565 = stablehlo.rsqrt %32564 : tensor<1x1x2048x1xbf16> loc(#loc2837)
      %32566 = stablehlo.divide %32565, %32564 : tensor<1x1x2048x1xbf16> loc(#loc2838)
      %32567 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32568 = stablehlo.broadcast_in_dim %32567, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2839)
      %32569 = stablehlo.multiply %32568, %32566 : tensor<1x1x2048x1xbf16> loc(#loc2839)
      %32570 = stablehlo.broadcast_in_dim %32565, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2839)
      %32571 = stablehlo.multiply %32561, %32570 : tensor<1x1x2048x2048xbf16> loc(#loc2839)
      %32572 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32573 = stablehlo.broadcast_in_dim %32572, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2840)
      %32574 = stablehlo.add %32573, %30649#0 : tensor<1x2048xbf16> loc(#loc2840)
      %32575 = stablehlo.broadcast_in_dim %32574, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2841)
      %32576 = stablehlo.broadcast_in_dim %32575, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2842)
      %32577 = stablehlo.multiply %32571, %32576 : tensor<1x1x2048x2048xbf16> loc(#loc2842)
      %32578 = stablehlo.broadcast_in_dim %30648#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2843)
      %32579 = stablehlo.broadcast_in_dim %32578, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2844)
      %32580 = stablehlo.add %32577, %32579 : tensor<1x1x2048x2048xbf16> loc(#loc2844)
      %32581 = func.call @_einsum_703(%32580, %30650#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2845)
      %32582 = "stablehlo.slice"(%32581) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2846)
      %32583 = stablehlo.reshape %32582 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2847)
      %32584 = "stablehlo.slice"(%32581) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2848)
      %32585 = stablehlo.reshape %32584 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2847)
      %32586 = "stablehlo.slice"(%32581) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2849)
      %32587 = stablehlo.reshape %32586 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2847)
      %32588 = arith.select %ptrue, %32583, %32583 : tensor<1x1x2048x32x64xbf16>
      %32589 = arith.select %ptrue, %32585, %32585 : tensor<1x1x2048x32x64xbf16>
      %32590 = arith.select %ptrue, %32587, %32587 : tensor<1x1x2048x32x64xbf16>
      %32591:2 = func.call @softplus_704(%30651#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2851)
      %32592 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %32593 = stablehlo.broadcast_in_dim %32592, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2852)
      %32594 = stablehlo.multiply %32593, %32591#0 : tensor<1x64xbf16> loc(#loc2852)
      %32595 = stablehlo.broadcast_in_dim %32594, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2853)
      %32596 = stablehlo.broadcast_in_dim %32595, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2854)
      %32597 = stablehlo.multiply %32588, %32596 : tensor<1x1x2048x32x64xbf16> loc(#loc2854)
      %32598 = func.call @_einsum_706(%32597, %32589) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2855)
      %32599 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32600 = stablehlo.broadcast_in_dim %32599, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2856)
      %32601 = stablehlo.divide %32598, %32600 : tensor<1x1x32x2048x2048xbf16> loc(#loc2856)
      %32602 = stablehlo.tanh %32601 : tensor<1x1x32x2048x2048xbf16> loc(#loc2857)
      %32603 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32604 = stablehlo.broadcast_in_dim %32603, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2858)
      %32605 = stablehlo.subtract %32604, %32602 : tensor<1x1x32x2048x2048xbf16> loc(#loc2858)
      %32606 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32607 = stablehlo.broadcast_in_dim %32606, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2859)
      %32608 = stablehlo.multiply %32607, %32602 : tensor<1x1x32x2048x2048xbf16> loc(#loc2859)
      %32609 = stablehlo.convert %32608 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2860)
      %32610 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32611 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32612 = stablehlo.multiply %32610, %32611 : tensor<f32> loc(#loc2861)
      %32613 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2862)
      %32614 = stablehlo.broadcast_in_dim %32612, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2863)
      %32615 = stablehlo.compare  GE, %32613, %32614,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2863)
      %32616 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32617:2 = func.call @_where_707(%32615, %32609, %32616) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2864)
      %32618 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2865)
      %32619 = stablehlo.reduce(%32617#0 init: %32618) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2865)
        stablehlo.return %39096 : tensor<f32> loc(#loc2865)
      } loc(#loc2865)
      %32620 = stablehlo.broadcast_in_dim %32619, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2866)
      %32621 = stablehlo.broadcast_in_dim %32620, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2867)
      %32622 = stablehlo.subtract %32617#0, %32621 : tensor<1x1x32x2048x2048xf32> loc(#loc2867)
      %32623 = stablehlo.exponential %32622 : tensor<1x1x32x2048x2048xf32> loc(#loc2868)
      %32624 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2869)
      %32625 = stablehlo.reduce(%32623 init: %32624) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2869)
        stablehlo.return %39096 : tensor<f32> loc(#loc2869)
      } loc(#loc2869)
      %32626 = stablehlo.broadcast_in_dim %32625, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2866)
      %32627 = stablehlo.broadcast_in_dim %32626, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2870)
      %32628 = stablehlo.divide %32623, %32627 : tensor<1x1x32x2048x2048xf32> loc(#loc2870)
      %32629 = stablehlo.multiply %32626, %32626 : tensor<1x1x32x2048x1xf32> loc(#loc2871)
      %32630 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2872)
      %32631 = stablehlo.broadcast_in_dim %32630, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2873)
      %32632 = stablehlo.divide %32631, %32629 : tensor<1x1x32x2048x1xf32> loc(#loc2870)
      %32633 = stablehlo.convert %32628 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2874)
      %32634 = func.call @_einsum_708(%32633, %32590) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2875)
      %32635 = arith.select %ptrue, %32634, %32634 : tensor<1x1x2048x32x64xbf16>
      %32636 = func.call @_einsum_709(%32635, %30652#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2876)
      %32637 = arith.select %ptrue, %32636, %32636 : tensor<1x1x2048x2048xbf16>
      %32638 = stablehlo.add %32637, %32479 : tensor<1x1x2048x2048xbf16> loc(#loc2878)
      %32639 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2879)
      %32640 = stablehlo.convert %32638 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2880)
      %32641 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2881)
      %32642 = stablehlo.reduce(%32640 init: %32641) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2881)
        stablehlo.return %39096 : tensor<f32> loc(#loc2881)
      } loc(#loc2881)
      %32643 = stablehlo.broadcast_in_dim %32642, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2882)
      %32644 = stablehlo.convert %32643 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2883)
      %32645 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32646 = stablehlo.broadcast_in_dim %32645, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2884)
      %32647 = stablehlo.divide %32644, %32646 : tensor<1x1x2048x1xbf16> loc(#loc2884)
      %32648 = stablehlo.broadcast_in_dim %32647, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2885)
      %32649 = stablehlo.subtract %32638, %32648 : tensor<1x1x2048x2048xbf16> loc(#loc2885)
      %32650 = stablehlo.multiply %32649, %32649 : tensor<1x1x2048x2048xbf16> loc(#loc2886)
      %32651 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32652 = stablehlo.broadcast_in_dim %32651, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2886)
      %32653 = stablehlo.multiply %32652, %32649 : tensor<1x1x2048x2048xbf16> loc(#loc2886)
      %32654 = stablehlo.convert %32650 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2880)
      %32655 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2881)
      %32656 = stablehlo.reduce(%32654 init: %32655) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2881)
        stablehlo.return %39096 : tensor<f32> loc(#loc2881)
      } loc(#loc2881)
      %32657 = stablehlo.broadcast_in_dim %32656, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2882)
      %32658 = stablehlo.convert %32657 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2883)
      %32659 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32660 = stablehlo.broadcast_in_dim %32659, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2884)
      %32661 = stablehlo.divide %32658, %32660 : tensor<1x1x2048x1xbf16> loc(#loc2884)
      %32662 = stablehlo.broadcast_in_dim %32647, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2885)
      %32663 = stablehlo.subtract %32638, %32662 : tensor<1x1x2048x2048xbf16> loc(#loc2885)
      %32664 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32665 = stablehlo.broadcast_in_dim %32664, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2887)
      %32666 = stablehlo.add %32661, %32665 : tensor<1x1x2048x1xbf16> loc(#loc2887)
      %32667 = stablehlo.rsqrt %32666 : tensor<1x1x2048x1xbf16> loc(#loc2888)
      %32668 = stablehlo.divide %32667, %32666 : tensor<1x1x2048x1xbf16> loc(#loc2889)
      %32669 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32670 = stablehlo.broadcast_in_dim %32669, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2890)
      %32671 = stablehlo.multiply %32670, %32668 : tensor<1x1x2048x1xbf16> loc(#loc2890)
      %32672 = stablehlo.broadcast_in_dim %32667, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2890)
      %32673 = stablehlo.multiply %32663, %32672 : tensor<1x1x2048x2048xbf16> loc(#loc2890)
      %32674 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32675 = stablehlo.broadcast_in_dim %32674, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2891)
      %32676 = stablehlo.add %32675, %30647#0 : tensor<1x2048xbf16> loc(#loc2891)
      %32677 = stablehlo.broadcast_in_dim %32676, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2892)
      %32678 = stablehlo.broadcast_in_dim %32677, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2893)
      %32679 = stablehlo.multiply %32673, %32678 : tensor<1x1x2048x2048xbf16> loc(#loc2893)
      %32680 = stablehlo.broadcast_in_dim %30646#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2894)
      %32681 = stablehlo.broadcast_in_dim %32680, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2895)
      %32682 = stablehlo.add %32679, %32681 : tensor<1x1x2048x2048xbf16> loc(#loc2895)
      %32683 = func.call @_einsum_710(%32682, %30643#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2896)
      %32684 = arith.select %ptrue, %32683, %32683 : tensor<1x1x2048x8192xbf16>
      %32685 = stablehlo.broadcast_in_dim %30642#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc2898)
      %32686 = stablehlo.broadcast_in_dim %32685, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2899)
      %32687 = stablehlo.add %32684, %32686 : tensor<1x1x2048x8192xbf16> loc(#loc2899)
      %32688 = stablehlo.multiply %32687, %32687 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32689 = stablehlo.multiply %32687, %32688 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32690 = stablehlo.multiply %32687, %32687 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32691 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %32692 = stablehlo.broadcast_in_dim %32691, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32693 = stablehlo.multiply %32692, %32690 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32694 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %32695 = stablehlo.broadcast_in_dim %32694, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32696 = stablehlo.multiply %32695, %32689 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32697 = stablehlo.add %32687, %32696 : tensor<1x1x2048x8192xbf16> loc(#loc2901)
      %32698 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %32699 = stablehlo.broadcast_in_dim %32698, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32700 = stablehlo.multiply %32699, %32697 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32701 = stablehlo.tanh %32700 : tensor<1x1x2048x8192xbf16> loc(#loc2902)
      %32702 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32703 = stablehlo.broadcast_in_dim %32702, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2903)
      %32704 = stablehlo.subtract %32703, %32701 : tensor<1x1x2048x8192xbf16> loc(#loc2903)
      %32705 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32706 = stablehlo.broadcast_in_dim %32705, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2901)
      %32707 = stablehlo.add %32706, %32701 : tensor<1x1x2048x8192xbf16> loc(#loc2901)
      %32708 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32709 = stablehlo.broadcast_in_dim %32708, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32710 = stablehlo.multiply %32709, %32707 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32711 = stablehlo.multiply %32687, %32710 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %32712 = stablehlo.convert %32639 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2904)
      %32713 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32714 = stablehlo.broadcast_in_dim %32713, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2905)
      %32715 = stablehlo.subtract %32714, %32712 : tensor<1x1x2048x1xf32> loc(#loc2905)
      %32716 = stablehlo.convert %32715 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2906)
      %32717 = stablehlo.broadcast_in_dim %32716, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2907)
      %32718 = stablehlo.multiply %32711, %32717 : tensor<1x1x2048x8192xbf16> loc(#loc2907)
      %32719 = func.call @_einsum_711(%32718, %30645#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2908)
      %32720 = arith.select %ptrue, %32719, %32719 : tensor<1x1x2048x2048xbf16>
      %32721 = stablehlo.broadcast_in_dim %30644#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2910)
      %32722 = stablehlo.broadcast_in_dim %32721, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2911)
      %32723 = stablehlo.add %32720, %32722 : tensor<1x1x2048x2048xbf16> loc(#loc2911)
      %32724 = stablehlo.convert %32639 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2904)
      %32725 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32726 = stablehlo.broadcast_in_dim %32725, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2905)
      %32727 = stablehlo.subtract %32726, %32724 : tensor<1x1x2048x1xf32> loc(#loc2905)
      %32728 = stablehlo.convert %32727 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2912)
      %32729 = stablehlo.broadcast_in_dim %32728, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2913)
      %32730 = stablehlo.multiply %32723, %32729 : tensor<1x1x2048x2048xbf16> loc(#loc2913)
      %32731 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32732 = stablehlo.broadcast_in_dim %32731, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2914)
      %32733 = stablehlo.multiply %32730, %32732 : tensor<1x1x2048x2048xbf16> loc(#loc2914)
      %32734 = stablehlo.add %32638, %32733 : tensor<1x1x2048x2048xbf16> loc(#loc2915)
      %32735 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2916)
      %32736 = stablehlo.convert %32734 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2917)
      %32737 = stablehlo.convert %32735 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc2918)
      %32738 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32739 = stablehlo.broadcast_in_dim %32738, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc2919)
      %32740 = stablehlo.subtract %32739, %32737 : tensor<1x1x2048x1xf32> loc(#loc2919)
      %32741 = stablehlo.convert %32740 : tensor<1x1x2048x1xf32> loc(#loc2920)
      %32742 = stablehlo.broadcast_in_dim %32741, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32743 = stablehlo.multiply %32736, %32742 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32744 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2922)
      %32745 = stablehlo.reduce(%32743 init: %32744) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2922)
        stablehlo.return %39096 : tensor<f32> loc(#loc2922)
      } loc(#loc2922)
      %32746 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32747 = stablehlo.broadcast_in_dim %32746, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc2923)
      %32748 = stablehlo.convert %32740 : tensor<1x1x2048x1xf32> loc(#loc2920)
      %32749 = stablehlo.broadcast_in_dim %32747, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2924)
      %32750 = stablehlo.broadcast_in_dim %32748, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32751 = stablehlo.multiply %32749, %32750 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32752 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2922)
      %32753 = stablehlo.reduce(%32751 init: %32752) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2922)
        stablehlo.return %39096 : tensor<f32> loc(#loc2922)
      } loc(#loc2922)
      %32754 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32755 = stablehlo.broadcast_in_dim %32754, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2925)
      %32756 = stablehlo.maximum %32755, %32753 : tensor<1xf32> loc(#loc2925)
      %32757 = stablehlo.divide %32745, %32756 : tensor<1xf32> loc(#loc2926)
      %32758 = stablehlo.broadcast_in_dim %32757, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2927)
      %32759 = stablehlo.broadcast_in_dim %32758, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2928)
      %32760 = stablehlo.subtract %32736, %32759 : tensor<1x1x2048x2048xf32> loc(#loc2928)
      %32761 = stablehlo.convert %32740 : tensor<1x1x2048x1xf32> loc(#loc2920)
      %32762 = stablehlo.broadcast_in_dim %32761, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32763 = stablehlo.multiply %32760, %32762 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32764 = stablehlo.multiply %32763, %32763 : tensor<1x1x2048x2048xf32> loc(#loc2929)
      %32765 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %32766 = stablehlo.broadcast_in_dim %32765, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2929)
      %32767 = stablehlo.multiply %32766, %32763 : tensor<1x1x2048x2048xf32> loc(#loc2929)
      %32768 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2922)
      %32769 = stablehlo.reduce(%32764 init: %32768) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2922)
        stablehlo.return %39096 : tensor<f32> loc(#loc2922)
      } loc(#loc2922)
      %32770 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32771 = stablehlo.broadcast_in_dim %32770, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2925)
      %32772 = stablehlo.maximum %32771, %32753 : tensor<1xf32> loc(#loc2925)
      %32773 = stablehlo.divide %32769, %32772 : tensor<1xf32> loc(#loc2926)
      %32774 = stablehlo.sqrt %32773 : tensor<1xf32> loc(#loc2930)
      %32775 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32776 = stablehlo.broadcast_in_dim %32775, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc2931)
      %32777 = stablehlo.divide %32776, %32774 : tensor<1xf32> loc(#loc2931)
      %32778 = stablehlo.convert %32740 : tensor<1x1x2048x1xf32> loc(#loc2920)
      %32779 = stablehlo.broadcast_in_dim %32778, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32780 = stablehlo.multiply %32736, %32779 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %32781 = stablehlo.abs %32780 : tensor<1x1x2048x2048xf32> loc(#loc2932)
      %32782 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %32783 = stablehlo.broadcast_in_dim %32782, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc2933)
      %32784 = stablehlo.compare  GE, %32780, %32783,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2933)
      %32785 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2934)
      %32786 = stablehlo.reduce(%32781 init: %32785) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2934)
        stablehlo.return %39096 : tensor<f32> loc(#loc2934)
      } loc(#loc2934)
      %32787 = stablehlo.reshape %32786 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc2935)
      %32788 = stablehlo.broadcast_in_dim %32787, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2936)
      %32789 = stablehlo.compare  EQ, %32781, %32788,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc2936)
      %32790 = stablehlo.convert %32789 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc2937)
      %32791 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2938)
      %32792 = stablehlo.reduce(%32790 init: %32791) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2938)
        stablehlo.return %39096 : tensor<f32> loc(#loc2938)
      } loc(#loc2938)
      %32793 = stablehlo.convert %32734 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2939)
      %32794 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2940)
      %32795 = stablehlo.reduce(%32793 init: %32794) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2940)
        stablehlo.return %39096 : tensor<f32> loc(#loc2940)
      } loc(#loc2940)
      %32796 = stablehlo.broadcast_in_dim %32795, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2941)
      %32797 = stablehlo.convert %32796 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2942)
      %32798 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32799 = stablehlo.broadcast_in_dim %32798, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2943)
      %32800 = stablehlo.divide %32797, %32799 : tensor<1x1x2048x1xbf16> loc(#loc2943)
      %32801 = stablehlo.broadcast_in_dim %32800, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2944)
      %32802 = stablehlo.subtract %32734, %32801 : tensor<1x1x2048x2048xbf16> loc(#loc2944)
      %32803 = stablehlo.multiply %32802, %32802 : tensor<1x1x2048x2048xbf16> loc(#loc2945)
      %32804 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32805 = stablehlo.broadcast_in_dim %32804, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2945)
      %32806 = stablehlo.multiply %32805, %32802 : tensor<1x1x2048x2048xbf16> loc(#loc2945)
      %32807 = stablehlo.convert %32803 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2939)
      %32808 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2940)
      %32809 = stablehlo.reduce(%32807 init: %32808) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2940)
        stablehlo.return %39096 : tensor<f32> loc(#loc2940)
      } loc(#loc2940)
      %32810 = stablehlo.broadcast_in_dim %32809, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2941)
      %32811 = stablehlo.convert %32810 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2942)
      %32812 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32813 = stablehlo.broadcast_in_dim %32812, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2943)
      %32814 = stablehlo.divide %32811, %32813 : tensor<1x1x2048x1xbf16> loc(#loc2943)
      %32815 = stablehlo.broadcast_in_dim %32800, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2944)
      %32816 = stablehlo.subtract %32734, %32815 : tensor<1x1x2048x2048xbf16> loc(#loc2944)
      %32817 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32818 = stablehlo.broadcast_in_dim %32817, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2946)
      %32819 = stablehlo.add %32814, %32818 : tensor<1x1x2048x1xbf16> loc(#loc2946)
      %32820 = stablehlo.rsqrt %32819 : tensor<1x1x2048x1xbf16> loc(#loc2947)
      %32821 = stablehlo.divide %32820, %32819 : tensor<1x1x2048x1xbf16> loc(#loc2948)
      %32822 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32823 = stablehlo.broadcast_in_dim %32822, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2949)
      %32824 = stablehlo.multiply %32823, %32821 : tensor<1x1x2048x1xbf16> loc(#loc2949)
      %32825 = stablehlo.broadcast_in_dim %32820, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2949)
      %32826 = stablehlo.multiply %32816, %32825 : tensor<1x1x2048x2048xbf16> loc(#loc2949)
      %32827 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32828 = stablehlo.broadcast_in_dim %32827, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc2950)
      %32829 = stablehlo.add %32828, %30660#0 : tensor<1x2048xbf16> loc(#loc2950)
      %32830 = stablehlo.broadcast_in_dim %32829, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2951)
      %32831 = stablehlo.broadcast_in_dim %32830, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2952)
      %32832 = stablehlo.multiply %32826, %32831 : tensor<1x1x2048x2048xbf16> loc(#loc2952)
      %32833 = stablehlo.broadcast_in_dim %30659#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc2953)
      %32834 = stablehlo.broadcast_in_dim %32833, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2954)
      %32835 = stablehlo.add %32832, %32834 : tensor<1x1x2048x2048xbf16> loc(#loc2954)
      %32836 = func.call @_einsum_712(%32835, %30661#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc2955)
      %32837 = "stablehlo.slice"(%32836) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2956)
      %32838 = stablehlo.reshape %32837 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2957)
      %32839 = "stablehlo.slice"(%32836) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2958)
      %32840 = stablehlo.reshape %32839 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2957)
      %32841 = "stablehlo.slice"(%32836) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc2959)
      %32842 = stablehlo.reshape %32841 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2957)
      %32843 = arith.select %ptrue, %32838, %32838 : tensor<1x1x2048x32x64xbf16>
      %32844 = arith.select %ptrue, %32840, %32840 : tensor<1x1x2048x32x64xbf16>
      %32845 = arith.select %ptrue, %32842, %32842 : tensor<1x1x2048x32x64xbf16>
      %32846:2 = func.call @softplus_713(%30662#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc2961)
      %32847 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %32848 = stablehlo.broadcast_in_dim %32847, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2962)
      %32849 = stablehlo.multiply %32848, %32846#0 : tensor<1x64xbf16> loc(#loc2962)
      %32850 = stablehlo.broadcast_in_dim %32849, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc2963)
      %32851 = stablehlo.broadcast_in_dim %32850, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2964)
      %32852 = stablehlo.multiply %32843, %32851 : tensor<1x1x2048x32x64xbf16> loc(#loc2964)
      %32853 = func.call @_einsum_715(%32852, %32844) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2965)
      %32854 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32855 = stablehlo.broadcast_in_dim %32854, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2966)
      %32856 = stablehlo.divide %32853, %32855 : tensor<1x1x32x2048x2048xbf16> loc(#loc2966)
      %32857 = stablehlo.tanh %32856 : tensor<1x1x32x2048x2048xbf16> loc(#loc2967)
      %32858 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32859 = stablehlo.broadcast_in_dim %32858, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2968)
      %32860 = stablehlo.subtract %32859, %32857 : tensor<1x1x32x2048x2048xbf16> loc(#loc2968)
      %32861 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %32862 = stablehlo.broadcast_in_dim %32861, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2969)
      %32863 = stablehlo.multiply %32862, %32857 : tensor<1x1x32x2048x2048xbf16> loc(#loc2969)
      %32864 = stablehlo.convert %32863 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2970)
      %32865 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32866 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %32867 = stablehlo.multiply %32865, %32866 : tensor<f32> loc(#loc2971)
      %32868 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2972)
      %32869 = stablehlo.broadcast_in_dim %32867, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc2973)
      %32870 = stablehlo.compare  GE, %32868, %32869,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc2973)
      %32871 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %32872:2 = func.call @_where_716(%32870, %32864, %32871) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc2974)
      %32873 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc2975)
      %32874 = stablehlo.reduce(%32872#0 init: %32873) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc2975)
        stablehlo.return %39096 : tensor<f32> loc(#loc2975)
      } loc(#loc2975)
      %32875 = stablehlo.broadcast_in_dim %32874, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2976)
      %32876 = stablehlo.broadcast_in_dim %32875, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2977)
      %32877 = stablehlo.subtract %32872#0, %32876 : tensor<1x1x32x2048x2048xf32> loc(#loc2977)
      %32878 = stablehlo.exponential %32877 : tensor<1x1x32x2048x2048xf32> loc(#loc2978)
      %32879 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2979)
      %32880 = stablehlo.reduce(%32878 init: %32879) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2979)
        stablehlo.return %39096 : tensor<f32> loc(#loc2979)
      } loc(#loc2979)
      %32881 = stablehlo.broadcast_in_dim %32880, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2976)
      %32882 = stablehlo.broadcast_in_dim %32881, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2980)
      %32883 = stablehlo.divide %32878, %32882 : tensor<1x1x32x2048x2048xf32> loc(#loc2980)
      %32884 = stablehlo.multiply %32881, %32881 : tensor<1x1x32x2048x1xf32> loc(#loc2981)
      %32885 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc2982)
      %32886 = stablehlo.broadcast_in_dim %32885, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc2983)
      %32887 = stablehlo.divide %32886, %32884 : tensor<1x1x32x2048x1xf32> loc(#loc2980)
      %32888 = stablehlo.convert %32883 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2984)
      %32889 = func.call @_einsum_717(%32888, %32845) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2985)
      %32890 = arith.select %ptrue, %32889, %32889 : tensor<1x1x2048x32x64xbf16>
      %32891 = func.call @_einsum_718(%32890, %30663#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2986)
      %32892 = arith.select %ptrue, %32891, %32891 : tensor<1x1x2048x2048xbf16>
      %32893 = stablehlo.add %32892, %32734 : tensor<1x1x2048x2048xbf16> loc(#loc2988)
      %32894 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc2989)
      %32895 = stablehlo.convert %32893 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2990)
      %32896 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2991)
      %32897 = stablehlo.reduce(%32895 init: %32896) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2991)
        stablehlo.return %39096 : tensor<f32> loc(#loc2991)
      } loc(#loc2991)
      %32898 = stablehlo.broadcast_in_dim %32897, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2992)
      %32899 = stablehlo.convert %32898 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2993)
      %32900 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32901 = stablehlo.broadcast_in_dim %32900, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2994)
      %32902 = stablehlo.divide %32899, %32901 : tensor<1x1x2048x1xbf16> loc(#loc2994)
      %32903 = stablehlo.broadcast_in_dim %32902, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2995)
      %32904 = stablehlo.subtract %32893, %32903 : tensor<1x1x2048x2048xbf16> loc(#loc2995)
      %32905 = stablehlo.multiply %32904, %32904 : tensor<1x1x2048x2048xbf16> loc(#loc2996)
      %32906 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %32907 = stablehlo.broadcast_in_dim %32906, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2996)
      %32908 = stablehlo.multiply %32907, %32904 : tensor<1x1x2048x2048xbf16> loc(#loc2996)
      %32909 = stablehlo.convert %32905 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc2990)
      %32910 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2991)
      %32911 = stablehlo.reduce(%32909 init: %32910) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2991)
        stablehlo.return %39096 : tensor<f32> loc(#loc2991)
      } loc(#loc2991)
      %32912 = stablehlo.broadcast_in_dim %32911, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc2992)
      %32913 = stablehlo.convert %32912 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc2993)
      %32914 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %32915 = stablehlo.broadcast_in_dim %32914, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2994)
      %32916 = stablehlo.divide %32913, %32915 : tensor<1x1x2048x1xbf16> loc(#loc2994)
      %32917 = stablehlo.broadcast_in_dim %32902, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2995)
      %32918 = stablehlo.subtract %32893, %32917 : tensor<1x1x2048x2048xbf16> loc(#loc2995)
      %32919 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %32920 = stablehlo.broadcast_in_dim %32919, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2997)
      %32921 = stablehlo.add %32916, %32920 : tensor<1x1x2048x1xbf16> loc(#loc2997)
      %32922 = stablehlo.rsqrt %32921 : tensor<1x1x2048x1xbf16> loc(#loc2998)
      %32923 = stablehlo.divide %32922, %32921 : tensor<1x1x2048x1xbf16> loc(#loc2999)
      %32924 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32925 = stablehlo.broadcast_in_dim %32924, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3000)
      %32926 = stablehlo.multiply %32925, %32923 : tensor<1x1x2048x1xbf16> loc(#loc3000)
      %32927 = stablehlo.broadcast_in_dim %32922, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3000)
      %32928 = stablehlo.multiply %32918, %32927 : tensor<1x1x2048x2048xbf16> loc(#loc3000)
      %32929 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32930 = stablehlo.broadcast_in_dim %32929, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3001)
      %32931 = stablehlo.add %32930, %30658#0 : tensor<1x2048xbf16> loc(#loc3001)
      %32932 = stablehlo.broadcast_in_dim %32931, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3002)
      %32933 = stablehlo.broadcast_in_dim %32932, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3003)
      %32934 = stablehlo.multiply %32928, %32933 : tensor<1x1x2048x2048xbf16> loc(#loc3003)
      %32935 = stablehlo.broadcast_in_dim %30657#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3004)
      %32936 = stablehlo.broadcast_in_dim %32935, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3005)
      %32937 = stablehlo.add %32934, %32936 : tensor<1x1x2048x2048xbf16> loc(#loc3005)
      %32938 = func.call @_einsum_719(%32937, %30654#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3006)
      %32939 = arith.select %ptrue, %32938, %32938 : tensor<1x1x2048x8192xbf16>
      %32940 = stablehlo.broadcast_in_dim %30653#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3008)
      %32941 = stablehlo.broadcast_in_dim %32940, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3009)
      %32942 = stablehlo.add %32939, %32941 : tensor<1x1x2048x8192xbf16> loc(#loc3009)
      %32943 = stablehlo.multiply %32942, %32942 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32944 = stablehlo.multiply %32942, %32943 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32945 = stablehlo.multiply %32942, %32942 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32946 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %32947 = stablehlo.broadcast_in_dim %32946, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32948 = stablehlo.multiply %32947, %32945 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32949 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %32950 = stablehlo.broadcast_in_dim %32949, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32951 = stablehlo.multiply %32950, %32944 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32952 = stablehlo.add %32942, %32951 : tensor<1x1x2048x8192xbf16> loc(#loc3011)
      %32953 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %32954 = stablehlo.broadcast_in_dim %32953, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32955 = stablehlo.multiply %32954, %32952 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32956 = stablehlo.tanh %32955 : tensor<1x1x2048x8192xbf16> loc(#loc3012)
      %32957 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32958 = stablehlo.broadcast_in_dim %32957, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3013)
      %32959 = stablehlo.subtract %32958, %32956 : tensor<1x1x2048x8192xbf16> loc(#loc3013)
      %32960 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32961 = stablehlo.broadcast_in_dim %32960, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3011)
      %32962 = stablehlo.add %32961, %32956 : tensor<1x1x2048x8192xbf16> loc(#loc3011)
      %32963 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %32964 = stablehlo.broadcast_in_dim %32963, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32965 = stablehlo.multiply %32964, %32962 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32966 = stablehlo.multiply %32942, %32965 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %32967 = stablehlo.convert %32894 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3014)
      %32968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32969 = stablehlo.broadcast_in_dim %32968, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3015)
      %32970 = stablehlo.subtract %32969, %32967 : tensor<1x1x2048x1xf32> loc(#loc3015)
      %32971 = stablehlo.convert %32970 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3016)
      %32972 = stablehlo.broadcast_in_dim %32971, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3017)
      %32973 = stablehlo.multiply %32966, %32972 : tensor<1x1x2048x8192xbf16> loc(#loc3017)
      %32974 = func.call @_einsum_720(%32973, %30656#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3018)
      %32975 = arith.select %ptrue, %32974, %32974 : tensor<1x1x2048x2048xbf16>
      %32976 = stablehlo.broadcast_in_dim %30655#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3020)
      %32977 = stablehlo.broadcast_in_dim %32976, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3021)
      %32978 = stablehlo.add %32975, %32977 : tensor<1x1x2048x2048xbf16> loc(#loc3021)
      %32979 = stablehlo.convert %32894 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3014)
      %32980 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32981 = stablehlo.broadcast_in_dim %32980, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3015)
      %32982 = stablehlo.subtract %32981, %32979 : tensor<1x1x2048x1xf32> loc(#loc3015)
      %32983 = stablehlo.convert %32982 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3022)
      %32984 = stablehlo.broadcast_in_dim %32983, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3023)
      %32985 = stablehlo.multiply %32978, %32984 : tensor<1x1x2048x2048xbf16> loc(#loc3023)
      %32986 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %32987 = stablehlo.broadcast_in_dim %32986, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3024)
      %32988 = stablehlo.multiply %32985, %32987 : tensor<1x1x2048x2048xbf16> loc(#loc3024)
      %32989 = stablehlo.add %32893, %32988 : tensor<1x1x2048x2048xbf16> loc(#loc3025)
      %32990 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3026)
      %32991 = stablehlo.convert %32989 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3027)
      %32992 = stablehlo.convert %32990 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3028)
      %32993 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %32994 = stablehlo.broadcast_in_dim %32993, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3029)
      %32995 = stablehlo.subtract %32994, %32992 : tensor<1x1x2048x1xf32> loc(#loc3029)
      %32996 = stablehlo.convert %32995 : tensor<1x1x2048x1xf32> loc(#loc3030)
      %32997 = stablehlo.broadcast_in_dim %32996, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %32998 = stablehlo.multiply %32991, %32997 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %32999 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3032)
      %33000 = stablehlo.reduce(%32998 init: %32999) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3032)
        stablehlo.return %39096 : tensor<f32> loc(#loc3032)
      } loc(#loc3032)
      %33001 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33002 = stablehlo.broadcast_in_dim %33001, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3033)
      %33003 = stablehlo.convert %32995 : tensor<1x1x2048x1xf32> loc(#loc3030)
      %33004 = stablehlo.broadcast_in_dim %33002, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3034)
      %33005 = stablehlo.broadcast_in_dim %33003, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %33006 = stablehlo.multiply %33004, %33005 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %33007 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3032)
      %33008 = stablehlo.reduce(%33006 init: %33007) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3032)
        stablehlo.return %39096 : tensor<f32> loc(#loc3032)
      } loc(#loc3032)
      %33009 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33010 = stablehlo.broadcast_in_dim %33009, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3035)
      %33011 = stablehlo.maximum %33010, %33008 : tensor<1xf32> loc(#loc3035)
      %33012 = stablehlo.divide %33000, %33011 : tensor<1xf32> loc(#loc3036)
      %33013 = stablehlo.broadcast_in_dim %33012, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3037)
      %33014 = stablehlo.broadcast_in_dim %33013, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3038)
      %33015 = stablehlo.subtract %32991, %33014 : tensor<1x1x2048x2048xf32> loc(#loc3038)
      %33016 = stablehlo.convert %32995 : tensor<1x1x2048x1xf32> loc(#loc3030)
      %33017 = stablehlo.broadcast_in_dim %33016, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %33018 = stablehlo.multiply %33015, %33017 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %33019 = stablehlo.multiply %33018, %33018 : tensor<1x1x2048x2048xf32> loc(#loc3039)
      %33020 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %33021 = stablehlo.broadcast_in_dim %33020, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3039)
      %33022 = stablehlo.multiply %33021, %33018 : tensor<1x1x2048x2048xf32> loc(#loc3039)
      %33023 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3032)
      %33024 = stablehlo.reduce(%33019 init: %33023) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3032)
        stablehlo.return %39096 : tensor<f32> loc(#loc3032)
      } loc(#loc3032)
      %33025 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33026 = stablehlo.broadcast_in_dim %33025, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3035)
      %33027 = stablehlo.maximum %33026, %33008 : tensor<1xf32> loc(#loc3035)
      %33028 = stablehlo.divide %33024, %33027 : tensor<1xf32> loc(#loc3036)
      %33029 = stablehlo.sqrt %33028 : tensor<1xf32> loc(#loc3040)
      %33030 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33031 = stablehlo.broadcast_in_dim %33030, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3041)
      %33032 = stablehlo.divide %33031, %33029 : tensor<1xf32> loc(#loc3041)
      %33033 = stablehlo.convert %32995 : tensor<1x1x2048x1xf32> loc(#loc3030)
      %33034 = stablehlo.broadcast_in_dim %33033, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %33035 = stablehlo.multiply %32991, %33034 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %33036 = stablehlo.abs %33035 : tensor<1x1x2048x2048xf32> loc(#loc3042)
      %33037 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %33038 = stablehlo.broadcast_in_dim %33037, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3043)
      %33039 = stablehlo.compare  GE, %33035, %33038,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3043)
      %33040 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3044)
      %33041 = stablehlo.reduce(%33036 init: %33040) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3044)
        stablehlo.return %39096 : tensor<f32> loc(#loc3044)
      } loc(#loc3044)
      %33042 = stablehlo.reshape %33041 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3045)
      %33043 = stablehlo.broadcast_in_dim %33042, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3046)
      %33044 = stablehlo.compare  EQ, %33036, %33043,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3046)
      %33045 = stablehlo.convert %33044 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3047)
      %33046 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3048)
      %33047 = stablehlo.reduce(%33045 init: %33046) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3048)
        stablehlo.return %39096 : tensor<f32> loc(#loc3048)
      } loc(#loc3048)
      %33048 = stablehlo.convert %32989 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3049)
      %33049 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3050)
      %33050 = stablehlo.reduce(%33048 init: %33049) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3050)
        stablehlo.return %39096 : tensor<f32> loc(#loc3050)
      } loc(#loc3050)
      %33051 = stablehlo.broadcast_in_dim %33050, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3051)
      %33052 = stablehlo.convert %33051 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3052)
      %33053 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33054 = stablehlo.broadcast_in_dim %33053, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3053)
      %33055 = stablehlo.divide %33052, %33054 : tensor<1x1x2048x1xbf16> loc(#loc3053)
      %33056 = stablehlo.broadcast_in_dim %33055, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3054)
      %33057 = stablehlo.subtract %32989, %33056 : tensor<1x1x2048x2048xbf16> loc(#loc3054)
      %33058 = stablehlo.multiply %33057, %33057 : tensor<1x1x2048x2048xbf16> loc(#loc3055)
      %33059 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33060 = stablehlo.broadcast_in_dim %33059, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3055)
      %33061 = stablehlo.multiply %33060, %33057 : tensor<1x1x2048x2048xbf16> loc(#loc3055)
      %33062 = stablehlo.convert %33058 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3049)
      %33063 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3050)
      %33064 = stablehlo.reduce(%33062 init: %33063) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3050)
        stablehlo.return %39096 : tensor<f32> loc(#loc3050)
      } loc(#loc3050)
      %33065 = stablehlo.broadcast_in_dim %33064, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3051)
      %33066 = stablehlo.convert %33065 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3052)
      %33067 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33068 = stablehlo.broadcast_in_dim %33067, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3053)
      %33069 = stablehlo.divide %33066, %33068 : tensor<1x1x2048x1xbf16> loc(#loc3053)
      %33070 = stablehlo.broadcast_in_dim %33055, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3054)
      %33071 = stablehlo.subtract %32989, %33070 : tensor<1x1x2048x2048xbf16> loc(#loc3054)
      %33072 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33073 = stablehlo.broadcast_in_dim %33072, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3056)
      %33074 = stablehlo.add %33069, %33073 : tensor<1x1x2048x1xbf16> loc(#loc3056)
      %33075 = stablehlo.rsqrt %33074 : tensor<1x1x2048x1xbf16> loc(#loc3057)
      %33076 = stablehlo.divide %33075, %33074 : tensor<1x1x2048x1xbf16> loc(#loc3058)
      %33077 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33078 = stablehlo.broadcast_in_dim %33077, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3059)
      %33079 = stablehlo.multiply %33078, %33076 : tensor<1x1x2048x1xbf16> loc(#loc3059)
      %33080 = stablehlo.broadcast_in_dim %33075, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3059)
      %33081 = stablehlo.multiply %33071, %33080 : tensor<1x1x2048x2048xbf16> loc(#loc3059)
      %33082 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33083 = stablehlo.broadcast_in_dim %33082, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3060)
      %33084 = stablehlo.add %33083, %30671#0 : tensor<1x2048xbf16> loc(#loc3060)
      %33085 = stablehlo.broadcast_in_dim %33084, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3061)
      %33086 = stablehlo.broadcast_in_dim %33085, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3062)
      %33087 = stablehlo.multiply %33081, %33086 : tensor<1x1x2048x2048xbf16> loc(#loc3062)
      %33088 = stablehlo.broadcast_in_dim %30670#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3063)
      %33089 = stablehlo.broadcast_in_dim %33088, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3064)
      %33090 = stablehlo.add %33087, %33089 : tensor<1x1x2048x2048xbf16> loc(#loc3064)
      %33091 = func.call @_einsum_721(%33090, %30672#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3065)
      %33092 = "stablehlo.slice"(%33091) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3066)
      %33093 = stablehlo.reshape %33092 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3067)
      %33094 = "stablehlo.slice"(%33091) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3068)
      %33095 = stablehlo.reshape %33094 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3067)
      %33096 = "stablehlo.slice"(%33091) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3069)
      %33097 = stablehlo.reshape %33096 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3067)
      %33098 = arith.select %ptrue, %33093, %33093 : tensor<1x1x2048x32x64xbf16>
      %33099 = arith.select %ptrue, %33095, %33095 : tensor<1x1x2048x32x64xbf16>
      %33100 = arith.select %ptrue, %33097, %33097 : tensor<1x1x2048x32x64xbf16>
      %33101:2 = func.call @softplus_722(%30673#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3071)
      %33102 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %33103 = stablehlo.broadcast_in_dim %33102, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3072)
      %33104 = stablehlo.multiply %33103, %33101#0 : tensor<1x64xbf16> loc(#loc3072)
      %33105 = stablehlo.broadcast_in_dim %33104, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3073)
      %33106 = stablehlo.broadcast_in_dim %33105, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3074)
      %33107 = stablehlo.multiply %33098, %33106 : tensor<1x1x2048x32x64xbf16> loc(#loc3074)
      %33108 = func.call @_einsum_724(%33107, %33099) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3075)
      %33109 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33110 = stablehlo.broadcast_in_dim %33109, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3076)
      %33111 = stablehlo.divide %33108, %33110 : tensor<1x1x32x2048x2048xbf16> loc(#loc3076)
      %33112 = stablehlo.tanh %33111 : tensor<1x1x32x2048x2048xbf16> loc(#loc3077)
      %33113 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33114 = stablehlo.broadcast_in_dim %33113, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3078)
      %33115 = stablehlo.subtract %33114, %33112 : tensor<1x1x32x2048x2048xbf16> loc(#loc3078)
      %33116 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33117 = stablehlo.broadcast_in_dim %33116, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3079)
      %33118 = stablehlo.multiply %33117, %33112 : tensor<1x1x32x2048x2048xbf16> loc(#loc3079)
      %33119 = stablehlo.convert %33118 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3080)
      %33120 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33121 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33122 = stablehlo.multiply %33120, %33121 : tensor<f32> loc(#loc3081)
      %33123 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3082)
      %33124 = stablehlo.broadcast_in_dim %33122, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3083)
      %33125 = stablehlo.compare  GE, %33123, %33124,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3083)
      %33126 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33127:2 = func.call @_where_725(%33125, %33119, %33126) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3084)
      %33128 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3085)
      %33129 = stablehlo.reduce(%33127#0 init: %33128) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3085)
        stablehlo.return %39096 : tensor<f32> loc(#loc3085)
      } loc(#loc3085)
      %33130 = stablehlo.broadcast_in_dim %33129, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3086)
      %33131 = stablehlo.broadcast_in_dim %33130, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3087)
      %33132 = stablehlo.subtract %33127#0, %33131 : tensor<1x1x32x2048x2048xf32> loc(#loc3087)
      %33133 = stablehlo.exponential %33132 : tensor<1x1x32x2048x2048xf32> loc(#loc3088)
      %33134 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3089)
      %33135 = stablehlo.reduce(%33133 init: %33134) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3089)
        stablehlo.return %39096 : tensor<f32> loc(#loc3089)
      } loc(#loc3089)
      %33136 = stablehlo.broadcast_in_dim %33135, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3086)
      %33137 = stablehlo.broadcast_in_dim %33136, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3090)
      %33138 = stablehlo.divide %33133, %33137 : tensor<1x1x32x2048x2048xf32> loc(#loc3090)
      %33139 = stablehlo.multiply %33136, %33136 : tensor<1x1x32x2048x1xf32> loc(#loc3091)
      %33140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3092)
      %33141 = stablehlo.broadcast_in_dim %33140, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3093)
      %33142 = stablehlo.divide %33141, %33139 : tensor<1x1x32x2048x1xf32> loc(#loc3090)
      %33143 = stablehlo.convert %33138 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3094)
      %33144 = func.call @_einsum_726(%33143, %33100) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3095)
      %33145 = arith.select %ptrue, %33144, %33144 : tensor<1x1x2048x32x64xbf16>
      %33146 = func.call @_einsum_727(%33145, %30674#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3096)
      %33147 = arith.select %ptrue, %33146, %33146 : tensor<1x1x2048x2048xbf16>
      %33148 = stablehlo.add %33147, %32989 : tensor<1x1x2048x2048xbf16> loc(#loc3098)
      %33149 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3099)
      %33150 = stablehlo.convert %33148 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3100)
      %33151 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3101)
      %33152 = stablehlo.reduce(%33150 init: %33151) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3101)
        stablehlo.return %39096 : tensor<f32> loc(#loc3101)
      } loc(#loc3101)
      %33153 = stablehlo.broadcast_in_dim %33152, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3102)
      %33154 = stablehlo.convert %33153 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3103)
      %33155 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33156 = stablehlo.broadcast_in_dim %33155, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3104)
      %33157 = stablehlo.divide %33154, %33156 : tensor<1x1x2048x1xbf16> loc(#loc3104)
      %33158 = stablehlo.broadcast_in_dim %33157, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3105)
      %33159 = stablehlo.subtract %33148, %33158 : tensor<1x1x2048x2048xbf16> loc(#loc3105)
      %33160 = stablehlo.multiply %33159, %33159 : tensor<1x1x2048x2048xbf16> loc(#loc3106)
      %33161 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33162 = stablehlo.broadcast_in_dim %33161, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3106)
      %33163 = stablehlo.multiply %33162, %33159 : tensor<1x1x2048x2048xbf16> loc(#loc3106)
      %33164 = stablehlo.convert %33160 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3100)
      %33165 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3101)
      %33166 = stablehlo.reduce(%33164 init: %33165) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3101)
        stablehlo.return %39096 : tensor<f32> loc(#loc3101)
      } loc(#loc3101)
      %33167 = stablehlo.broadcast_in_dim %33166, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3102)
      %33168 = stablehlo.convert %33167 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3103)
      %33169 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33170 = stablehlo.broadcast_in_dim %33169, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3104)
      %33171 = stablehlo.divide %33168, %33170 : tensor<1x1x2048x1xbf16> loc(#loc3104)
      %33172 = stablehlo.broadcast_in_dim %33157, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3105)
      %33173 = stablehlo.subtract %33148, %33172 : tensor<1x1x2048x2048xbf16> loc(#loc3105)
      %33174 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33175 = stablehlo.broadcast_in_dim %33174, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3107)
      %33176 = stablehlo.add %33171, %33175 : tensor<1x1x2048x1xbf16> loc(#loc3107)
      %33177 = stablehlo.rsqrt %33176 : tensor<1x1x2048x1xbf16> loc(#loc3108)
      %33178 = stablehlo.divide %33177, %33176 : tensor<1x1x2048x1xbf16> loc(#loc3109)
      %33179 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33180 = stablehlo.broadcast_in_dim %33179, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3110)
      %33181 = stablehlo.multiply %33180, %33178 : tensor<1x1x2048x1xbf16> loc(#loc3110)
      %33182 = stablehlo.broadcast_in_dim %33177, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3110)
      %33183 = stablehlo.multiply %33173, %33182 : tensor<1x1x2048x2048xbf16> loc(#loc3110)
      %33184 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33185 = stablehlo.broadcast_in_dim %33184, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3111)
      %33186 = stablehlo.add %33185, %30669#0 : tensor<1x2048xbf16> loc(#loc3111)
      %33187 = stablehlo.broadcast_in_dim %33186, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3112)
      %33188 = stablehlo.broadcast_in_dim %33187, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3113)
      %33189 = stablehlo.multiply %33183, %33188 : tensor<1x1x2048x2048xbf16> loc(#loc3113)
      %33190 = stablehlo.broadcast_in_dim %30668#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3114)
      %33191 = stablehlo.broadcast_in_dim %33190, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3115)
      %33192 = stablehlo.add %33189, %33191 : tensor<1x1x2048x2048xbf16> loc(#loc3115)
      %33193 = func.call @_einsum_728(%33192, %30665#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3116)
      %33194 = arith.select %ptrue, %33193, %33193 : tensor<1x1x2048x8192xbf16>
      %33195 = stablehlo.broadcast_in_dim %30664#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3118)
      %33196 = stablehlo.broadcast_in_dim %33195, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3119)
      %33197 = stablehlo.add %33194, %33196 : tensor<1x1x2048x8192xbf16> loc(#loc3119)
      %33198 = stablehlo.multiply %33197, %33197 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33199 = stablehlo.multiply %33197, %33198 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33200 = stablehlo.multiply %33197, %33197 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33201 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %33202 = stablehlo.broadcast_in_dim %33201, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33203 = stablehlo.multiply %33202, %33200 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33204 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %33205 = stablehlo.broadcast_in_dim %33204, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33206 = stablehlo.multiply %33205, %33199 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33207 = stablehlo.add %33197, %33206 : tensor<1x1x2048x8192xbf16> loc(#loc3121)
      %33208 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %33209 = stablehlo.broadcast_in_dim %33208, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33210 = stablehlo.multiply %33209, %33207 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33211 = stablehlo.tanh %33210 : tensor<1x1x2048x8192xbf16> loc(#loc3122)
      %33212 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33213 = stablehlo.broadcast_in_dim %33212, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3123)
      %33214 = stablehlo.subtract %33213, %33211 : tensor<1x1x2048x8192xbf16> loc(#loc3123)
      %33215 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33216 = stablehlo.broadcast_in_dim %33215, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3121)
      %33217 = stablehlo.add %33216, %33211 : tensor<1x1x2048x8192xbf16> loc(#loc3121)
      %33218 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33219 = stablehlo.broadcast_in_dim %33218, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33220 = stablehlo.multiply %33219, %33217 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33221 = stablehlo.multiply %33197, %33220 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %33222 = stablehlo.convert %33149 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3124)
      %33223 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33224 = stablehlo.broadcast_in_dim %33223, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3125)
      %33225 = stablehlo.subtract %33224, %33222 : tensor<1x1x2048x1xf32> loc(#loc3125)
      %33226 = stablehlo.convert %33225 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3126)
      %33227 = stablehlo.broadcast_in_dim %33226, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3127)
      %33228 = stablehlo.multiply %33221, %33227 : tensor<1x1x2048x8192xbf16> loc(#loc3127)
      %33229 = func.call @_einsum_729(%33228, %30667#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3128)
      %33230 = arith.select %ptrue, %33229, %33229 : tensor<1x1x2048x2048xbf16>
      %33231 = stablehlo.broadcast_in_dim %30666#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3130)
      %33232 = stablehlo.broadcast_in_dim %33231, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3131)
      %33233 = stablehlo.add %33230, %33232 : tensor<1x1x2048x2048xbf16> loc(#loc3131)
      %33234 = stablehlo.convert %33149 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3124)
      %33235 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33236 = stablehlo.broadcast_in_dim %33235, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3125)
      %33237 = stablehlo.subtract %33236, %33234 : tensor<1x1x2048x1xf32> loc(#loc3125)
      %33238 = stablehlo.convert %33237 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3132)
      %33239 = stablehlo.broadcast_in_dim %33238, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3133)
      %33240 = stablehlo.multiply %33233, %33239 : tensor<1x1x2048x2048xbf16> loc(#loc3133)
      %33241 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33242 = stablehlo.broadcast_in_dim %33241, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3134)
      %33243 = stablehlo.multiply %33240, %33242 : tensor<1x1x2048x2048xbf16> loc(#loc3134)
      %33244 = stablehlo.add %33148, %33243 : tensor<1x1x2048x2048xbf16> loc(#loc3135)
      %33245 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3136)
      %33246 = stablehlo.convert %33244 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3137)
      %33247 = stablehlo.convert %33245 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3138)
      %33248 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33249 = stablehlo.broadcast_in_dim %33248, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3139)
      %33250 = stablehlo.subtract %33249, %33247 : tensor<1x1x2048x1xf32> loc(#loc3139)
      %33251 = stablehlo.convert %33250 : tensor<1x1x2048x1xf32> loc(#loc3140)
      %33252 = stablehlo.broadcast_in_dim %33251, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33253 = stablehlo.multiply %33246, %33252 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33254 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3142)
      %33255 = stablehlo.reduce(%33253 init: %33254) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3142)
        stablehlo.return %39096 : tensor<f32> loc(#loc3142)
      } loc(#loc3142)
      %33256 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33257 = stablehlo.broadcast_in_dim %33256, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3143)
      %33258 = stablehlo.convert %33250 : tensor<1x1x2048x1xf32> loc(#loc3140)
      %33259 = stablehlo.broadcast_in_dim %33257, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3144)
      %33260 = stablehlo.broadcast_in_dim %33258, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33261 = stablehlo.multiply %33259, %33260 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33262 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3142)
      %33263 = stablehlo.reduce(%33261 init: %33262) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3142)
        stablehlo.return %39096 : tensor<f32> loc(#loc3142)
      } loc(#loc3142)
      %33264 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33265 = stablehlo.broadcast_in_dim %33264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3145)
      %33266 = stablehlo.maximum %33265, %33263 : tensor<1xf32> loc(#loc3145)
      %33267 = stablehlo.divide %33255, %33266 : tensor<1xf32> loc(#loc3146)
      %33268 = stablehlo.broadcast_in_dim %33267, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3147)
      %33269 = stablehlo.broadcast_in_dim %33268, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3148)
      %33270 = stablehlo.subtract %33246, %33269 : tensor<1x1x2048x2048xf32> loc(#loc3148)
      %33271 = stablehlo.convert %33250 : tensor<1x1x2048x1xf32> loc(#loc3140)
      %33272 = stablehlo.broadcast_in_dim %33271, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33273 = stablehlo.multiply %33270, %33272 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33274 = stablehlo.multiply %33273, %33273 : tensor<1x1x2048x2048xf32> loc(#loc3149)
      %33275 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %33276 = stablehlo.broadcast_in_dim %33275, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3149)
      %33277 = stablehlo.multiply %33276, %33273 : tensor<1x1x2048x2048xf32> loc(#loc3149)
      %33278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3142)
      %33279 = stablehlo.reduce(%33274 init: %33278) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3142)
        stablehlo.return %39096 : tensor<f32> loc(#loc3142)
      } loc(#loc3142)
      %33280 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33281 = stablehlo.broadcast_in_dim %33280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3145)
      %33282 = stablehlo.maximum %33281, %33263 : tensor<1xf32> loc(#loc3145)
      %33283 = stablehlo.divide %33279, %33282 : tensor<1xf32> loc(#loc3146)
      %33284 = stablehlo.sqrt %33283 : tensor<1xf32> loc(#loc3150)
      %33285 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33286 = stablehlo.broadcast_in_dim %33285, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3151)
      %33287 = stablehlo.divide %33286, %33284 : tensor<1xf32> loc(#loc3151)
      %33288 = stablehlo.convert %33250 : tensor<1x1x2048x1xf32> loc(#loc3140)
      %33289 = stablehlo.broadcast_in_dim %33288, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33290 = stablehlo.multiply %33246, %33289 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %33291 = stablehlo.abs %33290 : tensor<1x1x2048x2048xf32> loc(#loc3152)
      %33292 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %33293 = stablehlo.broadcast_in_dim %33292, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3153)
      %33294 = stablehlo.compare  GE, %33290, %33293,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3153)
      %33295 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3154)
      %33296 = stablehlo.reduce(%33291 init: %33295) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3154)
        stablehlo.return %39096 : tensor<f32> loc(#loc3154)
      } loc(#loc3154)
      %33297 = stablehlo.reshape %33296 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3155)
      %33298 = stablehlo.broadcast_in_dim %33297, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3156)
      %33299 = stablehlo.compare  EQ, %33291, %33298,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3156)
      %33300 = stablehlo.convert %33299 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3157)
      %33301 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3158)
      %33302 = stablehlo.reduce(%33300 init: %33301) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3158)
        stablehlo.return %39096 : tensor<f32> loc(#loc3158)
      } loc(#loc3158)
      %33303 = stablehlo.convert %33244 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3159)
      %33304 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3160)
      %33305 = stablehlo.reduce(%33303 init: %33304) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3160)
        stablehlo.return %39096 : tensor<f32> loc(#loc3160)
      } loc(#loc3160)
      %33306 = stablehlo.broadcast_in_dim %33305, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3161)
      %33307 = stablehlo.convert %33306 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3162)
      %33308 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33309 = stablehlo.broadcast_in_dim %33308, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3163)
      %33310 = stablehlo.divide %33307, %33309 : tensor<1x1x2048x1xbf16> loc(#loc3163)
      %33311 = stablehlo.broadcast_in_dim %33310, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3164)
      %33312 = stablehlo.subtract %33244, %33311 : tensor<1x1x2048x2048xbf16> loc(#loc3164)
      %33313 = stablehlo.multiply %33312, %33312 : tensor<1x1x2048x2048xbf16> loc(#loc3165)
      %33314 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33315 = stablehlo.broadcast_in_dim %33314, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3165)
      %33316 = stablehlo.multiply %33315, %33312 : tensor<1x1x2048x2048xbf16> loc(#loc3165)
      %33317 = stablehlo.convert %33313 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3159)
      %33318 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3160)
      %33319 = stablehlo.reduce(%33317 init: %33318) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3160)
        stablehlo.return %39096 : tensor<f32> loc(#loc3160)
      } loc(#loc3160)
      %33320 = stablehlo.broadcast_in_dim %33319, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3161)
      %33321 = stablehlo.convert %33320 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3162)
      %33322 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33323 = stablehlo.broadcast_in_dim %33322, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3163)
      %33324 = stablehlo.divide %33321, %33323 : tensor<1x1x2048x1xbf16> loc(#loc3163)
      %33325 = stablehlo.broadcast_in_dim %33310, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3164)
      %33326 = stablehlo.subtract %33244, %33325 : tensor<1x1x2048x2048xbf16> loc(#loc3164)
      %33327 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33328 = stablehlo.broadcast_in_dim %33327, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3166)
      %33329 = stablehlo.add %33324, %33328 : tensor<1x1x2048x1xbf16> loc(#loc3166)
      %33330 = stablehlo.rsqrt %33329 : tensor<1x1x2048x1xbf16> loc(#loc3167)
      %33331 = stablehlo.divide %33330, %33329 : tensor<1x1x2048x1xbf16> loc(#loc3168)
      %33332 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33333 = stablehlo.broadcast_in_dim %33332, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3169)
      %33334 = stablehlo.multiply %33333, %33331 : tensor<1x1x2048x1xbf16> loc(#loc3169)
      %33335 = stablehlo.broadcast_in_dim %33330, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3169)
      %33336 = stablehlo.multiply %33326, %33335 : tensor<1x1x2048x2048xbf16> loc(#loc3169)
      %33337 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33338 = stablehlo.broadcast_in_dim %33337, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3170)
      %33339 = stablehlo.add %33338, %30528#0 : tensor<1x2048xbf16> loc(#loc3170)
      %33340 = stablehlo.broadcast_in_dim %33339, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3171)
      %33341 = stablehlo.broadcast_in_dim %33340, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3172)
      %33342 = stablehlo.multiply %33336, %33341 : tensor<1x1x2048x2048xbf16> loc(#loc3172)
      %33343 = stablehlo.broadcast_in_dim %30527#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3173)
      %33344 = stablehlo.broadcast_in_dim %33343, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3174)
      %33345 = stablehlo.add %33342, %33344 : tensor<1x1x2048x2048xbf16> loc(#loc3174)
      %33346 = func.call @_einsum_730(%33345, %30529#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3175)
      %33347 = "stablehlo.slice"(%33346) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3176)
      %33348 = stablehlo.reshape %33347 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3177)
      %33349 = "stablehlo.slice"(%33346) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3178)
      %33350 = stablehlo.reshape %33349 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3177)
      %33351 = "stablehlo.slice"(%33346) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3179)
      %33352 = stablehlo.reshape %33351 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3177)
      %33353 = arith.select %ptrue, %33348, %33348 : tensor<1x1x2048x32x64xbf16>
      %33354 = arith.select %ptrue, %33350, %33350 : tensor<1x1x2048x32x64xbf16>
      %33355 = arith.select %ptrue, %33352, %33352 : tensor<1x1x2048x32x64xbf16>
      %33356:2 = func.call @softplus_731(%30530#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3181)
      %33357 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %33358 = stablehlo.broadcast_in_dim %33357, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3182)
      %33359 = stablehlo.multiply %33358, %33356#0 : tensor<1x64xbf16> loc(#loc3182)
      %33360 = stablehlo.broadcast_in_dim %33359, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3183)
      %33361 = stablehlo.broadcast_in_dim %33360, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3184)
      %33362 = stablehlo.multiply %33353, %33361 : tensor<1x1x2048x32x64xbf16> loc(#loc3184)
      %33363 = func.call @_einsum_733(%33362, %33354) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3185)
      %33364 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33365 = stablehlo.broadcast_in_dim %33364, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3186)
      %33366 = stablehlo.divide %33363, %33365 : tensor<1x1x32x2048x2048xbf16> loc(#loc3186)
      %33367 = stablehlo.tanh %33366 : tensor<1x1x32x2048x2048xbf16> loc(#loc3187)
      %33368 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33369 = stablehlo.broadcast_in_dim %33368, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3188)
      %33370 = stablehlo.subtract %33369, %33367 : tensor<1x1x32x2048x2048xbf16> loc(#loc3188)
      %33371 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33372 = stablehlo.broadcast_in_dim %33371, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3189)
      %33373 = stablehlo.multiply %33372, %33367 : tensor<1x1x32x2048x2048xbf16> loc(#loc3189)
      %33374 = stablehlo.convert %33373 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3190)
      %33375 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33376 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33377 = stablehlo.multiply %33375, %33376 : tensor<f32> loc(#loc3191)
      %33378 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3192)
      %33379 = stablehlo.broadcast_in_dim %33377, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3193)
      %33380 = stablehlo.compare  GE, %33378, %33379,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3193)
      %33381 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33382:2 = func.call @_where_734(%33380, %33374, %33381) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3194)
      %33383 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3195)
      %33384 = stablehlo.reduce(%33382#0 init: %33383) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3195)
        stablehlo.return %39096 : tensor<f32> loc(#loc3195)
      } loc(#loc3195)
      %33385 = stablehlo.broadcast_in_dim %33384, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3196)
      %33386 = stablehlo.broadcast_in_dim %33385, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3197)
      %33387 = stablehlo.subtract %33382#0, %33386 : tensor<1x1x32x2048x2048xf32> loc(#loc3197)
      %33388 = stablehlo.exponential %33387 : tensor<1x1x32x2048x2048xf32> loc(#loc3198)
      %33389 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3199)
      %33390 = stablehlo.reduce(%33388 init: %33389) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3199)
        stablehlo.return %39096 : tensor<f32> loc(#loc3199)
      } loc(#loc3199)
      %33391 = stablehlo.broadcast_in_dim %33390, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3196)
      %33392 = stablehlo.broadcast_in_dim %33391, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3200)
      %33393 = stablehlo.divide %33388, %33392 : tensor<1x1x32x2048x2048xf32> loc(#loc3200)
      %33394 = stablehlo.multiply %33391, %33391 : tensor<1x1x32x2048x1xf32> loc(#loc3201)
      %33395 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3202)
      %33396 = stablehlo.broadcast_in_dim %33395, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3203)
      %33397 = stablehlo.divide %33396, %33394 : tensor<1x1x32x2048x1xf32> loc(#loc3200)
      %33398 = stablehlo.convert %33393 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3204)
      %33399 = func.call @_einsum_735(%33398, %33355) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3205)
      %33400 = arith.select %ptrue, %33399, %33399 : tensor<1x1x2048x32x64xbf16>
      %33401 = func.call @_einsum_736(%33400, %30531#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3206)
      %33402 = arith.select %ptrue, %33401, %33401 : tensor<1x1x2048x2048xbf16>
      %33403 = stablehlo.add %33402, %33244 : tensor<1x1x2048x2048xbf16> loc(#loc3208)
      %33404 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3209)
      %33405 = stablehlo.convert %33403 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3210)
      %33406 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3211)
      %33407 = stablehlo.reduce(%33405 init: %33406) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3211)
        stablehlo.return %39096 : tensor<f32> loc(#loc3211)
      } loc(#loc3211)
      %33408 = stablehlo.broadcast_in_dim %33407, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3212)
      %33409 = stablehlo.convert %33408 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3213)
      %33410 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33411 = stablehlo.broadcast_in_dim %33410, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3214)
      %33412 = stablehlo.divide %33409, %33411 : tensor<1x1x2048x1xbf16> loc(#loc3214)
      %33413 = stablehlo.broadcast_in_dim %33412, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3215)
      %33414 = stablehlo.subtract %33403, %33413 : tensor<1x1x2048x2048xbf16> loc(#loc3215)
      %33415 = stablehlo.multiply %33414, %33414 : tensor<1x1x2048x2048xbf16> loc(#loc3216)
      %33416 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33417 = stablehlo.broadcast_in_dim %33416, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3216)
      %33418 = stablehlo.multiply %33417, %33414 : tensor<1x1x2048x2048xbf16> loc(#loc3216)
      %33419 = stablehlo.convert %33415 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3210)
      %33420 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3211)
      %33421 = stablehlo.reduce(%33419 init: %33420) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3211)
        stablehlo.return %39096 : tensor<f32> loc(#loc3211)
      } loc(#loc3211)
      %33422 = stablehlo.broadcast_in_dim %33421, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3212)
      %33423 = stablehlo.convert %33422 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3213)
      %33424 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33425 = stablehlo.broadcast_in_dim %33424, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3214)
      %33426 = stablehlo.divide %33423, %33425 : tensor<1x1x2048x1xbf16> loc(#loc3214)
      %33427 = stablehlo.broadcast_in_dim %33412, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3215)
      %33428 = stablehlo.subtract %33403, %33427 : tensor<1x1x2048x2048xbf16> loc(#loc3215)
      %33429 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33430 = stablehlo.broadcast_in_dim %33429, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3217)
      %33431 = stablehlo.add %33426, %33430 : tensor<1x1x2048x1xbf16> loc(#loc3217)
      %33432 = stablehlo.rsqrt %33431 : tensor<1x1x2048x1xbf16> loc(#loc3218)
      %33433 = stablehlo.divide %33432, %33431 : tensor<1x1x2048x1xbf16> loc(#loc3219)
      %33434 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33435 = stablehlo.broadcast_in_dim %33434, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3220)
      %33436 = stablehlo.multiply %33435, %33433 : tensor<1x1x2048x1xbf16> loc(#loc3220)
      %33437 = stablehlo.broadcast_in_dim %33432, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3220)
      %33438 = stablehlo.multiply %33428, %33437 : tensor<1x1x2048x2048xbf16> loc(#loc3220)
      %33439 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33440 = stablehlo.broadcast_in_dim %33439, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3221)
      %33441 = stablehlo.add %33440, %30526#0 : tensor<1x2048xbf16> loc(#loc3221)
      %33442 = stablehlo.broadcast_in_dim %33441, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3222)
      %33443 = stablehlo.broadcast_in_dim %33442, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3223)
      %33444 = stablehlo.multiply %33438, %33443 : tensor<1x1x2048x2048xbf16> loc(#loc3223)
      %33445 = stablehlo.broadcast_in_dim %30525#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3224)
      %33446 = stablehlo.broadcast_in_dim %33445, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3225)
      %33447 = stablehlo.add %33444, %33446 : tensor<1x1x2048x2048xbf16> loc(#loc3225)
      %33448 = func.call @_einsum_737(%33447, %30522#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3226)
      %33449 = arith.select %ptrue, %33448, %33448 : tensor<1x1x2048x8192xbf16>
      %33450 = stablehlo.broadcast_in_dim %30521#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3228)
      %33451 = stablehlo.broadcast_in_dim %33450, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3229)
      %33452 = stablehlo.add %33449, %33451 : tensor<1x1x2048x8192xbf16> loc(#loc3229)
      %33453 = stablehlo.multiply %33452, %33452 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33454 = stablehlo.multiply %33452, %33453 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33455 = stablehlo.multiply %33452, %33452 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33456 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %33457 = stablehlo.broadcast_in_dim %33456, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33458 = stablehlo.multiply %33457, %33455 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33459 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %33460 = stablehlo.broadcast_in_dim %33459, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33461 = stablehlo.multiply %33460, %33454 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33462 = stablehlo.add %33452, %33461 : tensor<1x1x2048x8192xbf16> loc(#loc3231)
      %33463 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %33464 = stablehlo.broadcast_in_dim %33463, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33465 = stablehlo.multiply %33464, %33462 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33466 = stablehlo.tanh %33465 : tensor<1x1x2048x8192xbf16> loc(#loc3232)
      %33467 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33468 = stablehlo.broadcast_in_dim %33467, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3233)
      %33469 = stablehlo.subtract %33468, %33466 : tensor<1x1x2048x8192xbf16> loc(#loc3233)
      %33470 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33471 = stablehlo.broadcast_in_dim %33470, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3231)
      %33472 = stablehlo.add %33471, %33466 : tensor<1x1x2048x8192xbf16> loc(#loc3231)
      %33473 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33474 = stablehlo.broadcast_in_dim %33473, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33475 = stablehlo.multiply %33474, %33472 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33476 = stablehlo.multiply %33452, %33475 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %33477 = stablehlo.convert %33404 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3234)
      %33478 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33479 = stablehlo.broadcast_in_dim %33478, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3235)
      %33480 = stablehlo.subtract %33479, %33477 : tensor<1x1x2048x1xf32> loc(#loc3235)
      %33481 = stablehlo.convert %33480 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3236)
      %33482 = stablehlo.broadcast_in_dim %33481, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3237)
      %33483 = stablehlo.multiply %33476, %33482 : tensor<1x1x2048x8192xbf16> loc(#loc3237)
      %33484 = func.call @_einsum_738(%33483, %30524#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3238)
      %33485 = arith.select %ptrue, %33484, %33484 : tensor<1x1x2048x2048xbf16>
      %33486 = stablehlo.broadcast_in_dim %30523#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3240)
      %33487 = stablehlo.broadcast_in_dim %33486, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3241)
      %33488 = stablehlo.add %33485, %33487 : tensor<1x1x2048x2048xbf16> loc(#loc3241)
      %33489 = stablehlo.convert %33404 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3234)
      %33490 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33491 = stablehlo.broadcast_in_dim %33490, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3235)
      %33492 = stablehlo.subtract %33491, %33489 : tensor<1x1x2048x1xf32> loc(#loc3235)
      %33493 = stablehlo.convert %33492 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3242)
      %33494 = stablehlo.broadcast_in_dim %33493, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3243)
      %33495 = stablehlo.multiply %33488, %33494 : tensor<1x1x2048x2048xbf16> loc(#loc3243)
      %33496 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33497 = stablehlo.broadcast_in_dim %33496, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3244)
      %33498 = stablehlo.multiply %33495, %33497 : tensor<1x1x2048x2048xbf16> loc(#loc3244)
      %33499 = stablehlo.add %33403, %33498 : tensor<1x1x2048x2048xbf16> loc(#loc3245)
      %33500 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3246)
      %33501 = stablehlo.convert %33499 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3247)
      %33502 = stablehlo.convert %33500 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3248)
      %33503 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33504 = stablehlo.broadcast_in_dim %33503, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3249)
      %33505 = stablehlo.subtract %33504, %33502 : tensor<1x1x2048x1xf32> loc(#loc3249)
      %33506 = stablehlo.convert %33505 : tensor<1x1x2048x1xf32> loc(#loc3250)
      %33507 = stablehlo.broadcast_in_dim %33506, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33508 = stablehlo.multiply %33501, %33507 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33509 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3252)
      %33510 = stablehlo.reduce(%33508 init: %33509) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3252)
        stablehlo.return %39096 : tensor<f32> loc(#loc3252)
      } loc(#loc3252)
      %33511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33512 = stablehlo.broadcast_in_dim %33511, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3253)
      %33513 = stablehlo.convert %33505 : tensor<1x1x2048x1xf32> loc(#loc3250)
      %33514 = stablehlo.broadcast_in_dim %33512, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3254)
      %33515 = stablehlo.broadcast_in_dim %33513, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33516 = stablehlo.multiply %33514, %33515 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33517 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3252)
      %33518 = stablehlo.reduce(%33516 init: %33517) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3252)
        stablehlo.return %39096 : tensor<f32> loc(#loc3252)
      } loc(#loc3252)
      %33519 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33520 = stablehlo.broadcast_in_dim %33519, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3255)
      %33521 = stablehlo.maximum %33520, %33518 : tensor<1xf32> loc(#loc3255)
      %33522 = stablehlo.divide %33510, %33521 : tensor<1xf32> loc(#loc3256)
      %33523 = stablehlo.broadcast_in_dim %33522, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3257)
      %33524 = stablehlo.broadcast_in_dim %33523, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3258)
      %33525 = stablehlo.subtract %33501, %33524 : tensor<1x1x2048x2048xf32> loc(#loc3258)
      %33526 = stablehlo.convert %33505 : tensor<1x1x2048x1xf32> loc(#loc3250)
      %33527 = stablehlo.broadcast_in_dim %33526, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33528 = stablehlo.multiply %33525, %33527 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33529 = stablehlo.multiply %33528, %33528 : tensor<1x1x2048x2048xf32> loc(#loc3259)
      %33530 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %33531 = stablehlo.broadcast_in_dim %33530, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3259)
      %33532 = stablehlo.multiply %33531, %33528 : tensor<1x1x2048x2048xf32> loc(#loc3259)
      %33533 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3252)
      %33534 = stablehlo.reduce(%33529 init: %33533) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3252)
        stablehlo.return %39096 : tensor<f32> loc(#loc3252)
      } loc(#loc3252)
      %33535 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33536 = stablehlo.broadcast_in_dim %33535, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3255)
      %33537 = stablehlo.maximum %33536, %33518 : tensor<1xf32> loc(#loc3255)
      %33538 = stablehlo.divide %33534, %33537 : tensor<1xf32> loc(#loc3256)
      %33539 = stablehlo.sqrt %33538 : tensor<1xf32> loc(#loc3260)
      %33540 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33541 = stablehlo.broadcast_in_dim %33540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3261)
      %33542 = stablehlo.divide %33541, %33539 : tensor<1xf32> loc(#loc3261)
      %33543 = stablehlo.convert %33505 : tensor<1x1x2048x1xf32> loc(#loc3250)
      %33544 = stablehlo.broadcast_in_dim %33543, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33545 = stablehlo.multiply %33501, %33544 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %33546 = stablehlo.abs %33545 : tensor<1x1x2048x2048xf32> loc(#loc3262)
      %33547 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %33548 = stablehlo.broadcast_in_dim %33547, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3263)
      %33549 = stablehlo.compare  GE, %33545, %33548,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3263)
      %33550 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3264)
      %33551 = stablehlo.reduce(%33546 init: %33550) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3264)
        stablehlo.return %39096 : tensor<f32> loc(#loc3264)
      } loc(#loc3264)
      %33552 = stablehlo.reshape %33551 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3265)
      %33553 = stablehlo.broadcast_in_dim %33552, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3266)
      %33554 = stablehlo.compare  EQ, %33546, %33553,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3266)
      %33555 = stablehlo.convert %33554 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3267)
      %33556 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3268)
      %33557 = stablehlo.reduce(%33555 init: %33556) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3268)
        stablehlo.return %39096 : tensor<f32> loc(#loc3268)
      } loc(#loc3268)
      %33558 = stablehlo.convert %33499 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3269)
      %33559 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3270)
      %33560 = stablehlo.reduce(%33558 init: %33559) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3270)
        stablehlo.return %39096 : tensor<f32> loc(#loc3270)
      } loc(#loc3270)
      %33561 = stablehlo.broadcast_in_dim %33560, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3271)
      %33562 = stablehlo.convert %33561 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3272)
      %33563 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33564 = stablehlo.broadcast_in_dim %33563, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3273)
      %33565 = stablehlo.divide %33562, %33564 : tensor<1x1x2048x1xbf16> loc(#loc3273)
      %33566 = stablehlo.broadcast_in_dim %33565, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3274)
      %33567 = stablehlo.subtract %33499, %33566 : tensor<1x1x2048x2048xbf16> loc(#loc3274)
      %33568 = stablehlo.multiply %33567, %33567 : tensor<1x1x2048x2048xbf16> loc(#loc3275)
      %33569 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33570 = stablehlo.broadcast_in_dim %33569, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3275)
      %33571 = stablehlo.multiply %33570, %33567 : tensor<1x1x2048x2048xbf16> loc(#loc3275)
      %33572 = stablehlo.convert %33568 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3269)
      %33573 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3270)
      %33574 = stablehlo.reduce(%33572 init: %33573) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3270)
        stablehlo.return %39096 : tensor<f32> loc(#loc3270)
      } loc(#loc3270)
      %33575 = stablehlo.broadcast_in_dim %33574, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3271)
      %33576 = stablehlo.convert %33575 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3272)
      %33577 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33578 = stablehlo.broadcast_in_dim %33577, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3273)
      %33579 = stablehlo.divide %33576, %33578 : tensor<1x1x2048x1xbf16> loc(#loc3273)
      %33580 = stablehlo.broadcast_in_dim %33565, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3274)
      %33581 = stablehlo.subtract %33499, %33580 : tensor<1x1x2048x2048xbf16> loc(#loc3274)
      %33582 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33583 = stablehlo.broadcast_in_dim %33582, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3276)
      %33584 = stablehlo.add %33579, %33583 : tensor<1x1x2048x1xbf16> loc(#loc3276)
      %33585 = stablehlo.rsqrt %33584 : tensor<1x1x2048x1xbf16> loc(#loc3277)
      %33586 = stablehlo.divide %33585, %33584 : tensor<1x1x2048x1xbf16> loc(#loc3278)
      %33587 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33588 = stablehlo.broadcast_in_dim %33587, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3279)
      %33589 = stablehlo.multiply %33588, %33586 : tensor<1x1x2048x1xbf16> loc(#loc3279)
      %33590 = stablehlo.broadcast_in_dim %33585, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3279)
      %33591 = stablehlo.multiply %33581, %33590 : tensor<1x1x2048x2048xbf16> loc(#loc3279)
      %33592 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33593 = stablehlo.broadcast_in_dim %33592, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3280)
      %33594 = stablehlo.add %33593, %30539#0 : tensor<1x2048xbf16> loc(#loc3280)
      %33595 = stablehlo.broadcast_in_dim %33594, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3281)
      %33596 = stablehlo.broadcast_in_dim %33595, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3282)
      %33597 = stablehlo.multiply %33591, %33596 : tensor<1x1x2048x2048xbf16> loc(#loc3282)
      %33598 = stablehlo.broadcast_in_dim %30538#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3283)
      %33599 = stablehlo.broadcast_in_dim %33598, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3284)
      %33600 = stablehlo.add %33597, %33599 : tensor<1x1x2048x2048xbf16> loc(#loc3284)
      %33601 = func.call @_einsum_739(%33600, %30540#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3285)
      %33602 = "stablehlo.slice"(%33601) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3286)
      %33603 = stablehlo.reshape %33602 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3287)
      %33604 = "stablehlo.slice"(%33601) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3288)
      %33605 = stablehlo.reshape %33604 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3287)
      %33606 = "stablehlo.slice"(%33601) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3289)
      %33607 = stablehlo.reshape %33606 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3287)
      %33608 = arith.select %ptrue, %33603, %33603 : tensor<1x1x2048x32x64xbf16>
      %33609 = arith.select %ptrue, %33605, %33605 : tensor<1x1x2048x32x64xbf16>
      %33610 = arith.select %ptrue, %33607, %33607 : tensor<1x1x2048x32x64xbf16>
      %33611:2 = func.call @softplus_740(%30541#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3291)
      %33612 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %33613 = stablehlo.broadcast_in_dim %33612, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3292)
      %33614 = stablehlo.multiply %33613, %33611#0 : tensor<1x64xbf16> loc(#loc3292)
      %33615 = stablehlo.broadcast_in_dim %33614, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3293)
      %33616 = stablehlo.broadcast_in_dim %33615, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3294)
      %33617 = stablehlo.multiply %33608, %33616 : tensor<1x1x2048x32x64xbf16> loc(#loc3294)
      %33618 = func.call @_einsum_742(%33617, %33609) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3295)
      %33619 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33620 = stablehlo.broadcast_in_dim %33619, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3296)
      %33621 = stablehlo.divide %33618, %33620 : tensor<1x1x32x2048x2048xbf16> loc(#loc3296)
      %33622 = stablehlo.tanh %33621 : tensor<1x1x32x2048x2048xbf16> loc(#loc3297)
      %33623 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33624 = stablehlo.broadcast_in_dim %33623, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3298)
      %33625 = stablehlo.subtract %33624, %33622 : tensor<1x1x32x2048x2048xbf16> loc(#loc3298)
      %33626 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33627 = stablehlo.broadcast_in_dim %33626, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3299)
      %33628 = stablehlo.multiply %33627, %33622 : tensor<1x1x32x2048x2048xbf16> loc(#loc3299)
      %33629 = stablehlo.convert %33628 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3300)
      %33630 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33631 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33632 = stablehlo.multiply %33630, %33631 : tensor<f32> loc(#loc3301)
      %33633 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3302)
      %33634 = stablehlo.broadcast_in_dim %33632, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3303)
      %33635 = stablehlo.compare  GE, %33633, %33634,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3303)
      %33636 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33637:2 = func.call @_where_743(%33635, %33629, %33636) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3304)
      %33638 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3305)
      %33639 = stablehlo.reduce(%33637#0 init: %33638) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3305)
        stablehlo.return %39096 : tensor<f32> loc(#loc3305)
      } loc(#loc3305)
      %33640 = stablehlo.broadcast_in_dim %33639, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3306)
      %33641 = stablehlo.broadcast_in_dim %33640, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3307)
      %33642 = stablehlo.subtract %33637#0, %33641 : tensor<1x1x32x2048x2048xf32> loc(#loc3307)
      %33643 = stablehlo.exponential %33642 : tensor<1x1x32x2048x2048xf32> loc(#loc3308)
      %33644 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3309)
      %33645 = stablehlo.reduce(%33643 init: %33644) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3309)
        stablehlo.return %39096 : tensor<f32> loc(#loc3309)
      } loc(#loc3309)
      %33646 = stablehlo.broadcast_in_dim %33645, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3306)
      %33647 = stablehlo.broadcast_in_dim %33646, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3310)
      %33648 = stablehlo.divide %33643, %33647 : tensor<1x1x32x2048x2048xf32> loc(#loc3310)
      %33649 = stablehlo.multiply %33646, %33646 : tensor<1x1x32x2048x1xf32> loc(#loc3311)
      %33650 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3312)
      %33651 = stablehlo.broadcast_in_dim %33650, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3313)
      %33652 = stablehlo.divide %33651, %33649 : tensor<1x1x32x2048x1xf32> loc(#loc3310)
      %33653 = stablehlo.convert %33648 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3314)
      %33654 = func.call @_einsum_744(%33653, %33610) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3315)
      %33655 = arith.select %ptrue, %33654, %33654 : tensor<1x1x2048x32x64xbf16>
      %33656 = func.call @_einsum_745(%33655, %30542#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3316)
      %33657 = arith.select %ptrue, %33656, %33656 : tensor<1x1x2048x2048xbf16>
      %33658 = stablehlo.add %33657, %33499 : tensor<1x1x2048x2048xbf16> loc(#loc3318)
      %33659 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3319)
      %33660 = stablehlo.convert %33658 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3320)
      %33661 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3321)
      %33662 = stablehlo.reduce(%33660 init: %33661) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3321)
        stablehlo.return %39096 : tensor<f32> loc(#loc3321)
      } loc(#loc3321)
      %33663 = stablehlo.broadcast_in_dim %33662, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3322)
      %33664 = stablehlo.convert %33663 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3323)
      %33665 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33666 = stablehlo.broadcast_in_dim %33665, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3324)
      %33667 = stablehlo.divide %33664, %33666 : tensor<1x1x2048x1xbf16> loc(#loc3324)
      %33668 = stablehlo.broadcast_in_dim %33667, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3325)
      %33669 = stablehlo.subtract %33658, %33668 : tensor<1x1x2048x2048xbf16> loc(#loc3325)
      %33670 = stablehlo.multiply %33669, %33669 : tensor<1x1x2048x2048xbf16> loc(#loc3326)
      %33671 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33672 = stablehlo.broadcast_in_dim %33671, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3326)
      %33673 = stablehlo.multiply %33672, %33669 : tensor<1x1x2048x2048xbf16> loc(#loc3326)
      %33674 = stablehlo.convert %33670 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3320)
      %33675 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3321)
      %33676 = stablehlo.reduce(%33674 init: %33675) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3321)
        stablehlo.return %39096 : tensor<f32> loc(#loc3321)
      } loc(#loc3321)
      %33677 = stablehlo.broadcast_in_dim %33676, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3322)
      %33678 = stablehlo.convert %33677 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3323)
      %33679 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33680 = stablehlo.broadcast_in_dim %33679, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3324)
      %33681 = stablehlo.divide %33678, %33680 : tensor<1x1x2048x1xbf16> loc(#loc3324)
      %33682 = stablehlo.broadcast_in_dim %33667, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3325)
      %33683 = stablehlo.subtract %33658, %33682 : tensor<1x1x2048x2048xbf16> loc(#loc3325)
      %33684 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33685 = stablehlo.broadcast_in_dim %33684, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3327)
      %33686 = stablehlo.add %33681, %33685 : tensor<1x1x2048x1xbf16> loc(#loc3327)
      %33687 = stablehlo.rsqrt %33686 : tensor<1x1x2048x1xbf16> loc(#loc3328)
      %33688 = stablehlo.divide %33687, %33686 : tensor<1x1x2048x1xbf16> loc(#loc3329)
      %33689 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33690 = stablehlo.broadcast_in_dim %33689, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3330)
      %33691 = stablehlo.multiply %33690, %33688 : tensor<1x1x2048x1xbf16> loc(#loc3330)
      %33692 = stablehlo.broadcast_in_dim %33687, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3330)
      %33693 = stablehlo.multiply %33683, %33692 : tensor<1x1x2048x2048xbf16> loc(#loc3330)
      %33694 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33695 = stablehlo.broadcast_in_dim %33694, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3331)
      %33696 = stablehlo.add %33695, %30537#0 : tensor<1x2048xbf16> loc(#loc3331)
      %33697 = stablehlo.broadcast_in_dim %33696, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3332)
      %33698 = stablehlo.broadcast_in_dim %33697, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3333)
      %33699 = stablehlo.multiply %33693, %33698 : tensor<1x1x2048x2048xbf16> loc(#loc3333)
      %33700 = stablehlo.broadcast_in_dim %30536#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3334)
      %33701 = stablehlo.broadcast_in_dim %33700, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3335)
      %33702 = stablehlo.add %33699, %33701 : tensor<1x1x2048x2048xbf16> loc(#loc3335)
      %33703 = func.call @_einsum_746(%33702, %30533#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3336)
      %33704 = arith.select %ptrue, %33703, %33703 : tensor<1x1x2048x8192xbf16>
      %33705 = stablehlo.broadcast_in_dim %30532#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3338)
      %33706 = stablehlo.broadcast_in_dim %33705, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3339)
      %33707 = stablehlo.add %33704, %33706 : tensor<1x1x2048x8192xbf16> loc(#loc3339)
      %33708 = stablehlo.multiply %33707, %33707 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33709 = stablehlo.multiply %33707, %33708 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33710 = stablehlo.multiply %33707, %33707 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33711 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %33712 = stablehlo.broadcast_in_dim %33711, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33713 = stablehlo.multiply %33712, %33710 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33714 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %33715 = stablehlo.broadcast_in_dim %33714, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33716 = stablehlo.multiply %33715, %33709 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33717 = stablehlo.add %33707, %33716 : tensor<1x1x2048x8192xbf16> loc(#loc3341)
      %33718 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %33719 = stablehlo.broadcast_in_dim %33718, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33720 = stablehlo.multiply %33719, %33717 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33721 = stablehlo.tanh %33720 : tensor<1x1x2048x8192xbf16> loc(#loc3342)
      %33722 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33723 = stablehlo.broadcast_in_dim %33722, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3343)
      %33724 = stablehlo.subtract %33723, %33721 : tensor<1x1x2048x8192xbf16> loc(#loc3343)
      %33725 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33726 = stablehlo.broadcast_in_dim %33725, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3341)
      %33727 = stablehlo.add %33726, %33721 : tensor<1x1x2048x8192xbf16> loc(#loc3341)
      %33728 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33729 = stablehlo.broadcast_in_dim %33728, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33730 = stablehlo.multiply %33729, %33727 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33731 = stablehlo.multiply %33707, %33730 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %33732 = stablehlo.convert %33659 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3344)
      %33733 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33734 = stablehlo.broadcast_in_dim %33733, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3345)
      %33735 = stablehlo.subtract %33734, %33732 : tensor<1x1x2048x1xf32> loc(#loc3345)
      %33736 = stablehlo.convert %33735 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3346)
      %33737 = stablehlo.broadcast_in_dim %33736, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3347)
      %33738 = stablehlo.multiply %33731, %33737 : tensor<1x1x2048x8192xbf16> loc(#loc3347)
      %33739 = func.call @_einsum_747(%33738, %30535#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3348)
      %33740 = arith.select %ptrue, %33739, %33739 : tensor<1x1x2048x2048xbf16>
      %33741 = stablehlo.broadcast_in_dim %30534#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3350)
      %33742 = stablehlo.broadcast_in_dim %33741, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3351)
      %33743 = stablehlo.add %33740, %33742 : tensor<1x1x2048x2048xbf16> loc(#loc3351)
      %33744 = stablehlo.convert %33659 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3344)
      %33745 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33746 = stablehlo.broadcast_in_dim %33745, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3345)
      %33747 = stablehlo.subtract %33746, %33744 : tensor<1x1x2048x1xf32> loc(#loc3345)
      %33748 = stablehlo.convert %33747 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3352)
      %33749 = stablehlo.broadcast_in_dim %33748, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3353)
      %33750 = stablehlo.multiply %33743, %33749 : tensor<1x1x2048x2048xbf16> loc(#loc3353)
      %33751 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33752 = stablehlo.broadcast_in_dim %33751, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3354)
      %33753 = stablehlo.multiply %33750, %33752 : tensor<1x1x2048x2048xbf16> loc(#loc3354)
      %33754 = stablehlo.add %33658, %33753 : tensor<1x1x2048x2048xbf16> loc(#loc3355)
      %33755 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3356)
      %33756 = stablehlo.convert %33754 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3357)
      %33757 = stablehlo.convert %33755 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3358)
      %33758 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33759 = stablehlo.broadcast_in_dim %33758, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3359)
      %33760 = stablehlo.subtract %33759, %33757 : tensor<1x1x2048x1xf32> loc(#loc3359)
      %33761 = stablehlo.convert %33760 : tensor<1x1x2048x1xf32> loc(#loc3360)
      %33762 = stablehlo.broadcast_in_dim %33761, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33763 = stablehlo.multiply %33756, %33762 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33764 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3362)
      %33765 = stablehlo.reduce(%33763 init: %33764) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3362)
        stablehlo.return %39096 : tensor<f32> loc(#loc3362)
      } loc(#loc3362)
      %33766 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33767 = stablehlo.broadcast_in_dim %33766, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3363)
      %33768 = stablehlo.convert %33760 : tensor<1x1x2048x1xf32> loc(#loc3360)
      %33769 = stablehlo.broadcast_in_dim %33767, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3364)
      %33770 = stablehlo.broadcast_in_dim %33768, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33771 = stablehlo.multiply %33769, %33770 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33772 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3362)
      %33773 = stablehlo.reduce(%33771 init: %33772) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3362)
        stablehlo.return %39096 : tensor<f32> loc(#loc3362)
      } loc(#loc3362)
      %33774 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33775 = stablehlo.broadcast_in_dim %33774, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3365)
      %33776 = stablehlo.maximum %33775, %33773 : tensor<1xf32> loc(#loc3365)
      %33777 = stablehlo.divide %33765, %33776 : tensor<1xf32> loc(#loc3366)
      %33778 = stablehlo.broadcast_in_dim %33777, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3367)
      %33779 = stablehlo.broadcast_in_dim %33778, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3368)
      %33780 = stablehlo.subtract %33756, %33779 : tensor<1x1x2048x2048xf32> loc(#loc3368)
      %33781 = stablehlo.convert %33760 : tensor<1x1x2048x1xf32> loc(#loc3360)
      %33782 = stablehlo.broadcast_in_dim %33781, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33783 = stablehlo.multiply %33780, %33782 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33784 = stablehlo.multiply %33783, %33783 : tensor<1x1x2048x2048xf32> loc(#loc3369)
      %33785 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %33786 = stablehlo.broadcast_in_dim %33785, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3369)
      %33787 = stablehlo.multiply %33786, %33783 : tensor<1x1x2048x2048xf32> loc(#loc3369)
      %33788 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3362)
      %33789 = stablehlo.reduce(%33784 init: %33788) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3362)
        stablehlo.return %39096 : tensor<f32> loc(#loc3362)
      } loc(#loc3362)
      %33790 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33791 = stablehlo.broadcast_in_dim %33790, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3365)
      %33792 = stablehlo.maximum %33791, %33773 : tensor<1xf32> loc(#loc3365)
      %33793 = stablehlo.divide %33789, %33792 : tensor<1xf32> loc(#loc3366)
      %33794 = stablehlo.sqrt %33793 : tensor<1xf32> loc(#loc3370)
      %33795 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33796 = stablehlo.broadcast_in_dim %33795, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3371)
      %33797 = stablehlo.divide %33796, %33794 : tensor<1xf32> loc(#loc3371)
      %33798 = stablehlo.convert %33760 : tensor<1x1x2048x1xf32> loc(#loc3360)
      %33799 = stablehlo.broadcast_in_dim %33798, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33800 = stablehlo.multiply %33756, %33799 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %33801 = stablehlo.abs %33800 : tensor<1x1x2048x2048xf32> loc(#loc3372)
      %33802 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %33803 = stablehlo.broadcast_in_dim %33802, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3373)
      %33804 = stablehlo.compare  GE, %33800, %33803,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3373)
      %33805 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3374)
      %33806 = stablehlo.reduce(%33801 init: %33805) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3374)
        stablehlo.return %39096 : tensor<f32> loc(#loc3374)
      } loc(#loc3374)
      %33807 = stablehlo.reshape %33806 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3375)
      %33808 = stablehlo.broadcast_in_dim %33807, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3376)
      %33809 = stablehlo.compare  EQ, %33801, %33808,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3376)
      %33810 = stablehlo.convert %33809 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3377)
      %33811 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3378)
      %33812 = stablehlo.reduce(%33810 init: %33811) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3378)
        stablehlo.return %39096 : tensor<f32> loc(#loc3378)
      } loc(#loc3378)
      %33813 = stablehlo.convert %33754 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3379)
      %33814 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3380)
      %33815 = stablehlo.reduce(%33813 init: %33814) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3380)
        stablehlo.return %39096 : tensor<f32> loc(#loc3380)
      } loc(#loc3380)
      %33816 = stablehlo.broadcast_in_dim %33815, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3381)
      %33817 = stablehlo.convert %33816 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3382)
      %33818 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33819 = stablehlo.broadcast_in_dim %33818, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3383)
      %33820 = stablehlo.divide %33817, %33819 : tensor<1x1x2048x1xbf16> loc(#loc3383)
      %33821 = stablehlo.broadcast_in_dim %33820, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3384)
      %33822 = stablehlo.subtract %33754, %33821 : tensor<1x1x2048x2048xbf16> loc(#loc3384)
      %33823 = stablehlo.multiply %33822, %33822 : tensor<1x1x2048x2048xbf16> loc(#loc3385)
      %33824 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33825 = stablehlo.broadcast_in_dim %33824, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3385)
      %33826 = stablehlo.multiply %33825, %33822 : tensor<1x1x2048x2048xbf16> loc(#loc3385)
      %33827 = stablehlo.convert %33823 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3379)
      %33828 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3380)
      %33829 = stablehlo.reduce(%33827 init: %33828) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3380)
        stablehlo.return %39096 : tensor<f32> loc(#loc3380)
      } loc(#loc3380)
      %33830 = stablehlo.broadcast_in_dim %33829, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3381)
      %33831 = stablehlo.convert %33830 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3382)
      %33832 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33833 = stablehlo.broadcast_in_dim %33832, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3383)
      %33834 = stablehlo.divide %33831, %33833 : tensor<1x1x2048x1xbf16> loc(#loc3383)
      %33835 = stablehlo.broadcast_in_dim %33820, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3384)
      %33836 = stablehlo.subtract %33754, %33835 : tensor<1x1x2048x2048xbf16> loc(#loc3384)
      %33837 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33838 = stablehlo.broadcast_in_dim %33837, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3386)
      %33839 = stablehlo.add %33834, %33838 : tensor<1x1x2048x1xbf16> loc(#loc3386)
      %33840 = stablehlo.rsqrt %33839 : tensor<1x1x2048x1xbf16> loc(#loc3387)
      %33841 = stablehlo.divide %33840, %33839 : tensor<1x1x2048x1xbf16> loc(#loc3388)
      %33842 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33843 = stablehlo.broadcast_in_dim %33842, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3389)
      %33844 = stablehlo.multiply %33843, %33841 : tensor<1x1x2048x1xbf16> loc(#loc3389)
      %33845 = stablehlo.broadcast_in_dim %33840, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3389)
      %33846 = stablehlo.multiply %33836, %33845 : tensor<1x1x2048x2048xbf16> loc(#loc3389)
      %33847 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33848 = stablehlo.broadcast_in_dim %33847, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3390)
      %33849 = stablehlo.add %33848, %30550#0 : tensor<1x2048xbf16> loc(#loc3390)
      %33850 = stablehlo.broadcast_in_dim %33849, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3391)
      %33851 = stablehlo.broadcast_in_dim %33850, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3392)
      %33852 = stablehlo.multiply %33846, %33851 : tensor<1x1x2048x2048xbf16> loc(#loc3392)
      %33853 = stablehlo.broadcast_in_dim %30549#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3393)
      %33854 = stablehlo.broadcast_in_dim %33853, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3394)
      %33855 = stablehlo.add %33852, %33854 : tensor<1x1x2048x2048xbf16> loc(#loc3394)
      %33856 = func.call @_einsum_748(%33855, %30551#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3395)
      %33857 = "stablehlo.slice"(%33856) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3396)
      %33858 = stablehlo.reshape %33857 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3397)
      %33859 = "stablehlo.slice"(%33856) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3398)
      %33860 = stablehlo.reshape %33859 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3397)
      %33861 = "stablehlo.slice"(%33856) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3399)
      %33862 = stablehlo.reshape %33861 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3397)
      %33863 = arith.select %ptrue, %33858, %33858 : tensor<1x1x2048x32x64xbf16>
      %33864 = arith.select %ptrue, %33860, %33860 : tensor<1x1x2048x32x64xbf16>
      %33865 = arith.select %ptrue, %33862, %33862 : tensor<1x1x2048x32x64xbf16>
      %33866:2 = func.call @softplus_749(%30552#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3401)
      %33867 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %33868 = stablehlo.broadcast_in_dim %33867, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3402)
      %33869 = stablehlo.multiply %33868, %33866#0 : tensor<1x64xbf16> loc(#loc3402)
      %33870 = stablehlo.broadcast_in_dim %33869, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3403)
      %33871 = stablehlo.broadcast_in_dim %33870, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3404)
      %33872 = stablehlo.multiply %33863, %33871 : tensor<1x1x2048x32x64xbf16> loc(#loc3404)
      %33873 = func.call @_einsum_751(%33872, %33864) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3405)
      %33874 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33875 = stablehlo.broadcast_in_dim %33874, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3406)
      %33876 = stablehlo.divide %33873, %33875 : tensor<1x1x32x2048x2048xbf16> loc(#loc3406)
      %33877 = stablehlo.tanh %33876 : tensor<1x1x32x2048x2048xbf16> loc(#loc3407)
      %33878 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33879 = stablehlo.broadcast_in_dim %33878, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3408)
      %33880 = stablehlo.subtract %33879, %33877 : tensor<1x1x32x2048x2048xbf16> loc(#loc3408)
      %33881 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %33882 = stablehlo.broadcast_in_dim %33881, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3409)
      %33883 = stablehlo.multiply %33882, %33877 : tensor<1x1x32x2048x2048xbf16> loc(#loc3409)
      %33884 = stablehlo.convert %33883 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3410)
      %33885 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33886 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %33887 = stablehlo.multiply %33885, %33886 : tensor<f32> loc(#loc3411)
      %33888 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3412)
      %33889 = stablehlo.broadcast_in_dim %33887, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3413)
      %33890 = stablehlo.compare  GE, %33888, %33889,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3413)
      %33891 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %33892:2 = func.call @_where_752(%33890, %33884, %33891) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3414)
      %33893 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3415)
      %33894 = stablehlo.reduce(%33892#0 init: %33893) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3415)
        stablehlo.return %39096 : tensor<f32> loc(#loc3415)
      } loc(#loc3415)
      %33895 = stablehlo.broadcast_in_dim %33894, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3416)
      %33896 = stablehlo.broadcast_in_dim %33895, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3417)
      %33897 = stablehlo.subtract %33892#0, %33896 : tensor<1x1x32x2048x2048xf32> loc(#loc3417)
      %33898 = stablehlo.exponential %33897 : tensor<1x1x32x2048x2048xf32> loc(#loc3418)
      %33899 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3419)
      %33900 = stablehlo.reduce(%33898 init: %33899) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3419)
        stablehlo.return %39096 : tensor<f32> loc(#loc3419)
      } loc(#loc3419)
      %33901 = stablehlo.broadcast_in_dim %33900, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3416)
      %33902 = stablehlo.broadcast_in_dim %33901, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3420)
      %33903 = stablehlo.divide %33898, %33902 : tensor<1x1x32x2048x2048xf32> loc(#loc3420)
      %33904 = stablehlo.multiply %33901, %33901 : tensor<1x1x32x2048x1xf32> loc(#loc3421)
      %33905 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3422)
      %33906 = stablehlo.broadcast_in_dim %33905, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3423)
      %33907 = stablehlo.divide %33906, %33904 : tensor<1x1x32x2048x1xf32> loc(#loc3420)
      %33908 = stablehlo.convert %33903 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3424)
      %33909 = func.call @_einsum_753(%33908, %33865) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3425)
      %33910 = arith.select %ptrue, %33909, %33909 : tensor<1x1x2048x32x64xbf16>
      %33911 = func.call @_einsum_754(%33910, %30553#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3426)
      %33912 = arith.select %ptrue, %33911, %33911 : tensor<1x1x2048x2048xbf16>
      %33913 = stablehlo.add %33912, %33754 : tensor<1x1x2048x2048xbf16> loc(#loc3428)
      %33914 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3429)
      %33915 = stablehlo.convert %33913 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3430)
      %33916 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3431)
      %33917 = stablehlo.reduce(%33915 init: %33916) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3431)
        stablehlo.return %39096 : tensor<f32> loc(#loc3431)
      } loc(#loc3431)
      %33918 = stablehlo.broadcast_in_dim %33917, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3432)
      %33919 = stablehlo.convert %33918 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3433)
      %33920 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33921 = stablehlo.broadcast_in_dim %33920, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3434)
      %33922 = stablehlo.divide %33919, %33921 : tensor<1x1x2048x1xbf16> loc(#loc3434)
      %33923 = stablehlo.broadcast_in_dim %33922, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3435)
      %33924 = stablehlo.subtract %33913, %33923 : tensor<1x1x2048x2048xbf16> loc(#loc3435)
      %33925 = stablehlo.multiply %33924, %33924 : tensor<1x1x2048x2048xbf16> loc(#loc3436)
      %33926 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %33927 = stablehlo.broadcast_in_dim %33926, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3436)
      %33928 = stablehlo.multiply %33927, %33924 : tensor<1x1x2048x2048xbf16> loc(#loc3436)
      %33929 = stablehlo.convert %33925 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3430)
      %33930 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3431)
      %33931 = stablehlo.reduce(%33929 init: %33930) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3431)
        stablehlo.return %39096 : tensor<f32> loc(#loc3431)
      } loc(#loc3431)
      %33932 = stablehlo.broadcast_in_dim %33931, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3432)
      %33933 = stablehlo.convert %33932 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3433)
      %33934 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %33935 = stablehlo.broadcast_in_dim %33934, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3434)
      %33936 = stablehlo.divide %33933, %33935 : tensor<1x1x2048x1xbf16> loc(#loc3434)
      %33937 = stablehlo.broadcast_in_dim %33922, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3435)
      %33938 = stablehlo.subtract %33913, %33937 : tensor<1x1x2048x2048xbf16> loc(#loc3435)
      %33939 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %33940 = stablehlo.broadcast_in_dim %33939, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3437)
      %33941 = stablehlo.add %33936, %33940 : tensor<1x1x2048x1xbf16> loc(#loc3437)
      %33942 = stablehlo.rsqrt %33941 : tensor<1x1x2048x1xbf16> loc(#loc3438)
      %33943 = stablehlo.divide %33942, %33941 : tensor<1x1x2048x1xbf16> loc(#loc3439)
      %33944 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33945 = stablehlo.broadcast_in_dim %33944, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3440)
      %33946 = stablehlo.multiply %33945, %33943 : tensor<1x1x2048x1xbf16> loc(#loc3440)
      %33947 = stablehlo.broadcast_in_dim %33942, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3440)
      %33948 = stablehlo.multiply %33938, %33947 : tensor<1x1x2048x2048xbf16> loc(#loc3440)
      %33949 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33950 = stablehlo.broadcast_in_dim %33949, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3441)
      %33951 = stablehlo.add %33950, %30548#0 : tensor<1x2048xbf16> loc(#loc3441)
      %33952 = stablehlo.broadcast_in_dim %33951, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3442)
      %33953 = stablehlo.broadcast_in_dim %33952, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3443)
      %33954 = stablehlo.multiply %33948, %33953 : tensor<1x1x2048x2048xbf16> loc(#loc3443)
      %33955 = stablehlo.broadcast_in_dim %30547#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3444)
      %33956 = stablehlo.broadcast_in_dim %33955, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3445)
      %33957 = stablehlo.add %33954, %33956 : tensor<1x1x2048x2048xbf16> loc(#loc3445)
      %33958 = func.call @_einsum_755(%33957, %30544#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3446)
      %33959 = arith.select %ptrue, %33958, %33958 : tensor<1x1x2048x8192xbf16>
      %33960 = stablehlo.broadcast_in_dim %30543#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3448)
      %33961 = stablehlo.broadcast_in_dim %33960, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3449)
      %33962 = stablehlo.add %33959, %33961 : tensor<1x1x2048x8192xbf16> loc(#loc3449)
      %33963 = stablehlo.multiply %33962, %33962 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33964 = stablehlo.multiply %33962, %33963 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33965 = stablehlo.multiply %33962, %33962 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33966 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %33967 = stablehlo.broadcast_in_dim %33966, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33968 = stablehlo.multiply %33967, %33965 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33969 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %33970 = stablehlo.broadcast_in_dim %33969, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33971 = stablehlo.multiply %33970, %33964 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33972 = stablehlo.add %33962, %33971 : tensor<1x1x2048x8192xbf16> loc(#loc3451)
      %33973 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %33974 = stablehlo.broadcast_in_dim %33973, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33975 = stablehlo.multiply %33974, %33972 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33976 = stablehlo.tanh %33975 : tensor<1x1x2048x8192xbf16> loc(#loc3452)
      %33977 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33978 = stablehlo.broadcast_in_dim %33977, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3453)
      %33979 = stablehlo.subtract %33978, %33976 : tensor<1x1x2048x8192xbf16> loc(#loc3453)
      %33980 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %33981 = stablehlo.broadcast_in_dim %33980, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3451)
      %33982 = stablehlo.add %33981, %33976 : tensor<1x1x2048x8192xbf16> loc(#loc3451)
      %33983 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %33984 = stablehlo.broadcast_in_dim %33983, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33985 = stablehlo.multiply %33984, %33982 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33986 = stablehlo.multiply %33962, %33985 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %33987 = stablehlo.convert %33914 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3454)
      %33988 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %33989 = stablehlo.broadcast_in_dim %33988, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3455)
      %33990 = stablehlo.subtract %33989, %33987 : tensor<1x1x2048x1xf32> loc(#loc3455)
      %33991 = stablehlo.convert %33990 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3456)
      %33992 = stablehlo.broadcast_in_dim %33991, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3457)
      %33993 = stablehlo.multiply %33986, %33992 : tensor<1x1x2048x8192xbf16> loc(#loc3457)
      %33994 = func.call @_einsum_756(%33993, %30546#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3458)
      %33995 = arith.select %ptrue, %33994, %33994 : tensor<1x1x2048x2048xbf16>
      %33996 = stablehlo.broadcast_in_dim %30545#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3460)
      %33997 = stablehlo.broadcast_in_dim %33996, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3461)
      %33998 = stablehlo.add %33995, %33997 : tensor<1x1x2048x2048xbf16> loc(#loc3461)
      %33999 = stablehlo.convert %33914 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3454)
      %34000 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34001 = stablehlo.broadcast_in_dim %34000, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3455)
      %34002 = stablehlo.subtract %34001, %33999 : tensor<1x1x2048x1xf32> loc(#loc3455)
      %34003 = stablehlo.convert %34002 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3462)
      %34004 = stablehlo.broadcast_in_dim %34003, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3463)
      %34005 = stablehlo.multiply %33998, %34004 : tensor<1x1x2048x2048xbf16> loc(#loc3463)
      %34006 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34007 = stablehlo.broadcast_in_dim %34006, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3464)
      %34008 = stablehlo.multiply %34005, %34007 : tensor<1x1x2048x2048xbf16> loc(#loc3464)
      %34009 = stablehlo.add %33913, %34008 : tensor<1x1x2048x2048xbf16> loc(#loc3465)
      %34010 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3466)
      %34011 = stablehlo.convert %34009 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3467)
      %34012 = stablehlo.convert %34010 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3468)
      %34013 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34014 = stablehlo.broadcast_in_dim %34013, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3469)
      %34015 = stablehlo.subtract %34014, %34012 : tensor<1x1x2048x1xf32> loc(#loc3469)
      %34016 = stablehlo.convert %34015 : tensor<1x1x2048x1xf32> loc(#loc3470)
      %34017 = stablehlo.broadcast_in_dim %34016, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34018 = stablehlo.multiply %34011, %34017 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34019 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3472)
      %34020 = stablehlo.reduce(%34018 init: %34019) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3472)
        stablehlo.return %39096 : tensor<f32> loc(#loc3472)
      } loc(#loc3472)
      %34021 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34022 = stablehlo.broadcast_in_dim %34021, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3473)
      %34023 = stablehlo.convert %34015 : tensor<1x1x2048x1xf32> loc(#loc3470)
      %34024 = stablehlo.broadcast_in_dim %34022, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3474)
      %34025 = stablehlo.broadcast_in_dim %34023, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34026 = stablehlo.multiply %34024, %34025 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34027 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3472)
      %34028 = stablehlo.reduce(%34026 init: %34027) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3472)
        stablehlo.return %39096 : tensor<f32> loc(#loc3472)
      } loc(#loc3472)
      %34029 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34030 = stablehlo.broadcast_in_dim %34029, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3475)
      %34031 = stablehlo.maximum %34030, %34028 : tensor<1xf32> loc(#loc3475)
      %34032 = stablehlo.divide %34020, %34031 : tensor<1xf32> loc(#loc3476)
      %34033 = stablehlo.broadcast_in_dim %34032, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3477)
      %34034 = stablehlo.broadcast_in_dim %34033, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3478)
      %34035 = stablehlo.subtract %34011, %34034 : tensor<1x1x2048x2048xf32> loc(#loc3478)
      %34036 = stablehlo.convert %34015 : tensor<1x1x2048x1xf32> loc(#loc3470)
      %34037 = stablehlo.broadcast_in_dim %34036, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34038 = stablehlo.multiply %34035, %34037 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34039 = stablehlo.multiply %34038, %34038 : tensor<1x1x2048x2048xf32> loc(#loc3479)
      %34040 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %34041 = stablehlo.broadcast_in_dim %34040, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3479)
      %34042 = stablehlo.multiply %34041, %34038 : tensor<1x1x2048x2048xf32> loc(#loc3479)
      %34043 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3472)
      %34044 = stablehlo.reduce(%34039 init: %34043) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3472)
        stablehlo.return %39096 : tensor<f32> loc(#loc3472)
      } loc(#loc3472)
      %34045 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34046 = stablehlo.broadcast_in_dim %34045, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3475)
      %34047 = stablehlo.maximum %34046, %34028 : tensor<1xf32> loc(#loc3475)
      %34048 = stablehlo.divide %34044, %34047 : tensor<1xf32> loc(#loc3476)
      %34049 = stablehlo.sqrt %34048 : tensor<1xf32> loc(#loc3480)
      %34050 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %34051 = stablehlo.broadcast_in_dim %34050, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3481)
      %34052 = stablehlo.divide %34051, %34049 : tensor<1xf32> loc(#loc3481)
      %34053 = stablehlo.convert %34015 : tensor<1x1x2048x1xf32> loc(#loc3470)
      %34054 = stablehlo.broadcast_in_dim %34053, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34055 = stablehlo.multiply %34011, %34054 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %34056 = stablehlo.abs %34055 : tensor<1x1x2048x2048xf32> loc(#loc3482)
      %34057 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34058 = stablehlo.broadcast_in_dim %34057, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3483)
      %34059 = stablehlo.compare  GE, %34055, %34058,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3483)
      %34060 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3484)
      %34061 = stablehlo.reduce(%34056 init: %34060) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3484)
        stablehlo.return %39096 : tensor<f32> loc(#loc3484)
      } loc(#loc3484)
      %34062 = stablehlo.reshape %34061 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3485)
      %34063 = stablehlo.broadcast_in_dim %34062, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3486)
      %34064 = stablehlo.compare  EQ, %34056, %34063,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3486)
      %34065 = stablehlo.convert %34064 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3487)
      %34066 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3488)
      %34067 = stablehlo.reduce(%34065 init: %34066) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3488)
        stablehlo.return %39096 : tensor<f32> loc(#loc3488)
      } loc(#loc3488)
      %34068 = stablehlo.convert %34009 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3489)
      %34069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3490)
      %34070 = stablehlo.reduce(%34068 init: %34069) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3490)
        stablehlo.return %39096 : tensor<f32> loc(#loc3490)
      } loc(#loc3490)
      %34071 = stablehlo.broadcast_in_dim %34070, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3491)
      %34072 = stablehlo.convert %34071 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3492)
      %34073 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34074 = stablehlo.broadcast_in_dim %34073, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3493)
      %34075 = stablehlo.divide %34072, %34074 : tensor<1x1x2048x1xbf16> loc(#loc3493)
      %34076 = stablehlo.broadcast_in_dim %34075, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3494)
      %34077 = stablehlo.subtract %34009, %34076 : tensor<1x1x2048x2048xbf16> loc(#loc3494)
      %34078 = stablehlo.multiply %34077, %34077 : tensor<1x1x2048x2048xbf16> loc(#loc3495)
      %34079 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %34080 = stablehlo.broadcast_in_dim %34079, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3495)
      %34081 = stablehlo.multiply %34080, %34077 : tensor<1x1x2048x2048xbf16> loc(#loc3495)
      %34082 = stablehlo.convert %34078 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3489)
      %34083 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3490)
      %34084 = stablehlo.reduce(%34082 init: %34083) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3490)
        stablehlo.return %39096 : tensor<f32> loc(#loc3490)
      } loc(#loc3490)
      %34085 = stablehlo.broadcast_in_dim %34084, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3491)
      %34086 = stablehlo.convert %34085 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3492)
      %34087 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34088 = stablehlo.broadcast_in_dim %34087, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3493)
      %34089 = stablehlo.divide %34086, %34088 : tensor<1x1x2048x1xbf16> loc(#loc3493)
      %34090 = stablehlo.broadcast_in_dim %34075, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3494)
      %34091 = stablehlo.subtract %34009, %34090 : tensor<1x1x2048x2048xbf16> loc(#loc3494)
      %34092 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %34093 = stablehlo.broadcast_in_dim %34092, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3496)
      %34094 = stablehlo.add %34089, %34093 : tensor<1x1x2048x1xbf16> loc(#loc3496)
      %34095 = stablehlo.rsqrt %34094 : tensor<1x1x2048x1xbf16> loc(#loc3497)
      %34096 = stablehlo.divide %34095, %34094 : tensor<1x1x2048x1xbf16> loc(#loc3498)
      %34097 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34098 = stablehlo.broadcast_in_dim %34097, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3499)
      %34099 = stablehlo.multiply %34098, %34096 : tensor<1x1x2048x1xbf16> loc(#loc3499)
      %34100 = stablehlo.broadcast_in_dim %34095, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3499)
      %34101 = stablehlo.multiply %34091, %34100 : tensor<1x1x2048x2048xbf16> loc(#loc3499)
      %34102 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34103 = stablehlo.broadcast_in_dim %34102, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3500)
      %34104 = stablehlo.add %34103, %30561#0 : tensor<1x2048xbf16> loc(#loc3500)
      %34105 = stablehlo.broadcast_in_dim %34104, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3501)
      %34106 = stablehlo.broadcast_in_dim %34105, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3502)
      %34107 = stablehlo.multiply %34101, %34106 : tensor<1x1x2048x2048xbf16> loc(#loc3502)
      %34108 = stablehlo.broadcast_in_dim %30560#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3503)
      %34109 = stablehlo.broadcast_in_dim %34108, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3504)
      %34110 = stablehlo.add %34107, %34109 : tensor<1x1x2048x2048xbf16> loc(#loc3504)
      %34111 = func.call @_einsum_757(%34110, %30562#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3505)
      %34112 = "stablehlo.slice"(%34111) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3506)
      %34113 = stablehlo.reshape %34112 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3507)
      %34114 = "stablehlo.slice"(%34111) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3508)
      %34115 = stablehlo.reshape %34114 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3507)
      %34116 = "stablehlo.slice"(%34111) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3509)
      %34117 = stablehlo.reshape %34116 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3507)
      %34118 = arith.select %ptrue, %34113, %34113 : tensor<1x1x2048x32x64xbf16>
      %34119 = arith.select %ptrue, %34115, %34115 : tensor<1x1x2048x32x64xbf16>
      %34120 = arith.select %ptrue, %34117, %34117 : tensor<1x1x2048x32x64xbf16>
      %34121:2 = func.call @softplus_758(%30563#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3511)
      %34122 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %34123 = stablehlo.broadcast_in_dim %34122, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3512)
      %34124 = stablehlo.multiply %34123, %34121#0 : tensor<1x64xbf16> loc(#loc3512)
      %34125 = stablehlo.broadcast_in_dim %34124, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3513)
      %34126 = stablehlo.broadcast_in_dim %34125, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3514)
      %34127 = stablehlo.multiply %34118, %34126 : tensor<1x1x2048x32x64xbf16> loc(#loc3514)
      %34128 = func.call @_einsum_760(%34127, %34119) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3515)
      %34129 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %34130 = stablehlo.broadcast_in_dim %34129, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3516)
      %34131 = stablehlo.divide %34128, %34130 : tensor<1x1x32x2048x2048xbf16> loc(#loc3516)
      %34132 = stablehlo.tanh %34131 : tensor<1x1x32x2048x2048xbf16> loc(#loc3517)
      %34133 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34134 = stablehlo.broadcast_in_dim %34133, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3518)
      %34135 = stablehlo.subtract %34134, %34132 : tensor<1x1x32x2048x2048xbf16> loc(#loc3518)
      %34136 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %34137 = stablehlo.broadcast_in_dim %34136, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3519)
      %34138 = stablehlo.multiply %34137, %34132 : tensor<1x1x32x2048x2048xbf16> loc(#loc3519)
      %34139 = stablehlo.convert %34138 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3520)
      %34140 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %34141 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %34142 = stablehlo.multiply %34140, %34141 : tensor<f32> loc(#loc3521)
      %34143 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3522)
      %34144 = stablehlo.broadcast_in_dim %34142, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3523)
      %34145 = stablehlo.compare  GE, %34143, %34144,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3523)
      %34146 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %34147:2 = func.call @_where_761(%34145, %34139, %34146) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3524)
      %34148 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3525)
      %34149 = stablehlo.reduce(%34147#0 init: %34148) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3525)
        stablehlo.return %39096 : tensor<f32> loc(#loc3525)
      } loc(#loc3525)
      %34150 = stablehlo.broadcast_in_dim %34149, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3526)
      %34151 = stablehlo.broadcast_in_dim %34150, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3527)
      %34152 = stablehlo.subtract %34147#0, %34151 : tensor<1x1x32x2048x2048xf32> loc(#loc3527)
      %34153 = stablehlo.exponential %34152 : tensor<1x1x32x2048x2048xf32> loc(#loc3528)
      %34154 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3529)
      %34155 = stablehlo.reduce(%34153 init: %34154) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3529)
        stablehlo.return %39096 : tensor<f32> loc(#loc3529)
      } loc(#loc3529)
      %34156 = stablehlo.broadcast_in_dim %34155, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3526)
      %34157 = stablehlo.broadcast_in_dim %34156, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3530)
      %34158 = stablehlo.divide %34153, %34157 : tensor<1x1x32x2048x2048xf32> loc(#loc3530)
      %34159 = stablehlo.multiply %34156, %34156 : tensor<1x1x32x2048x1xf32> loc(#loc3531)
      %34160 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3532)
      %34161 = stablehlo.broadcast_in_dim %34160, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3533)
      %34162 = stablehlo.divide %34161, %34159 : tensor<1x1x32x2048x1xf32> loc(#loc3530)
      %34163 = stablehlo.convert %34158 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3534)
      %34164 = func.call @_einsum_762(%34163, %34120) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3535)
      %34165 = arith.select %ptrue, %34164, %34164 : tensor<1x1x2048x32x64xbf16>
      %34166 = func.call @_einsum_763(%34165, %30564#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3536)
      %34167 = arith.select %ptrue, %34166, %34166 : tensor<1x1x2048x2048xbf16>
      %34168 = stablehlo.add %34167, %34009 : tensor<1x1x2048x2048xbf16> loc(#loc3538)
      %34169 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3539)
      %34170 = stablehlo.convert %34168 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3540)
      %34171 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3541)
      %34172 = stablehlo.reduce(%34170 init: %34171) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3541)
        stablehlo.return %39096 : tensor<f32> loc(#loc3541)
      } loc(#loc3541)
      %34173 = stablehlo.broadcast_in_dim %34172, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3542)
      %34174 = stablehlo.convert %34173 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3543)
      %34175 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34176 = stablehlo.broadcast_in_dim %34175, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3544)
      %34177 = stablehlo.divide %34174, %34176 : tensor<1x1x2048x1xbf16> loc(#loc3544)
      %34178 = stablehlo.broadcast_in_dim %34177, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3545)
      %34179 = stablehlo.subtract %34168, %34178 : tensor<1x1x2048x2048xbf16> loc(#loc3545)
      %34180 = stablehlo.multiply %34179, %34179 : tensor<1x1x2048x2048xbf16> loc(#loc3546)
      %34181 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %34182 = stablehlo.broadcast_in_dim %34181, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3546)
      %34183 = stablehlo.multiply %34182, %34179 : tensor<1x1x2048x2048xbf16> loc(#loc3546)
      %34184 = stablehlo.convert %34180 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3540)
      %34185 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3541)
      %34186 = stablehlo.reduce(%34184 init: %34185) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3541)
        stablehlo.return %39096 : tensor<f32> loc(#loc3541)
      } loc(#loc3541)
      %34187 = stablehlo.broadcast_in_dim %34186, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3542)
      %34188 = stablehlo.convert %34187 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3543)
      %34189 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34190 = stablehlo.broadcast_in_dim %34189, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3544)
      %34191 = stablehlo.divide %34188, %34190 : tensor<1x1x2048x1xbf16> loc(#loc3544)
      %34192 = stablehlo.broadcast_in_dim %34177, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3545)
      %34193 = stablehlo.subtract %34168, %34192 : tensor<1x1x2048x2048xbf16> loc(#loc3545)
      %34194 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %34195 = stablehlo.broadcast_in_dim %34194, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3547)
      %34196 = stablehlo.add %34191, %34195 : tensor<1x1x2048x1xbf16> loc(#loc3547)
      %34197 = stablehlo.rsqrt %34196 : tensor<1x1x2048x1xbf16> loc(#loc3548)
      %34198 = stablehlo.divide %34197, %34196 : tensor<1x1x2048x1xbf16> loc(#loc3549)
      %34199 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34200 = stablehlo.broadcast_in_dim %34199, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3550)
      %34201 = stablehlo.multiply %34200, %34198 : tensor<1x1x2048x1xbf16> loc(#loc3550)
      %34202 = stablehlo.broadcast_in_dim %34197, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3550)
      %34203 = stablehlo.multiply %34193, %34202 : tensor<1x1x2048x2048xbf16> loc(#loc3550)
      %34204 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34205 = stablehlo.broadcast_in_dim %34204, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3551)
      %34206 = stablehlo.add %34205, %30559#0 : tensor<1x2048xbf16> loc(#loc3551)
      %34207 = stablehlo.broadcast_in_dim %34206, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3552)
      %34208 = stablehlo.broadcast_in_dim %34207, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3553)
      %34209 = stablehlo.multiply %34203, %34208 : tensor<1x1x2048x2048xbf16> loc(#loc3553)
      %34210 = stablehlo.broadcast_in_dim %30558#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3554)
      %34211 = stablehlo.broadcast_in_dim %34210, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3555)
      %34212 = stablehlo.add %34209, %34211 : tensor<1x1x2048x2048xbf16> loc(#loc3555)
      %34213 = func.call @_einsum_764(%34212, %30555#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3556)
      %34214 = arith.select %ptrue, %34213, %34213 : tensor<1x1x2048x8192xbf16>
      %34215 = stablehlo.broadcast_in_dim %30554#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3558)
      %34216 = stablehlo.broadcast_in_dim %34215, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3559)
      %34217 = stablehlo.add %34214, %34216 : tensor<1x1x2048x8192xbf16> loc(#loc3559)
      %34218 = stablehlo.multiply %34217, %34217 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34219 = stablehlo.multiply %34217, %34218 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34220 = stablehlo.multiply %34217, %34217 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34221 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %34222 = stablehlo.broadcast_in_dim %34221, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34223 = stablehlo.multiply %34222, %34220 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34224 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %34225 = stablehlo.broadcast_in_dim %34224, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34226 = stablehlo.multiply %34225, %34219 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34227 = stablehlo.add %34217, %34226 : tensor<1x1x2048x8192xbf16> loc(#loc3561)
      %34228 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %34229 = stablehlo.broadcast_in_dim %34228, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34230 = stablehlo.multiply %34229, %34227 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34231 = stablehlo.tanh %34230 : tensor<1x1x2048x8192xbf16> loc(#loc3562)
      %34232 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34233 = stablehlo.broadcast_in_dim %34232, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3563)
      %34234 = stablehlo.subtract %34233, %34231 : tensor<1x1x2048x8192xbf16> loc(#loc3563)
      %34235 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34236 = stablehlo.broadcast_in_dim %34235, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3561)
      %34237 = stablehlo.add %34236, %34231 : tensor<1x1x2048x8192xbf16> loc(#loc3561)
      %34238 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34239 = stablehlo.broadcast_in_dim %34238, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34240 = stablehlo.multiply %34239, %34237 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34241 = stablehlo.multiply %34217, %34240 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %34242 = stablehlo.convert %34169 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3564)
      %34243 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34244 = stablehlo.broadcast_in_dim %34243, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3565)
      %34245 = stablehlo.subtract %34244, %34242 : tensor<1x1x2048x1xf32> loc(#loc3565)
      %34246 = stablehlo.convert %34245 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3566)
      %34247 = stablehlo.broadcast_in_dim %34246, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3567)
      %34248 = stablehlo.multiply %34241, %34247 : tensor<1x1x2048x8192xbf16> loc(#loc3567)
      %34249 = func.call @_einsum_765(%34248, %30557#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3568)
      %34250 = arith.select %ptrue, %34249, %34249 : tensor<1x1x2048x2048xbf16>
      %34251 = stablehlo.broadcast_in_dim %30556#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3570)
      %34252 = stablehlo.broadcast_in_dim %34251, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3571)
      %34253 = stablehlo.add %34250, %34252 : tensor<1x1x2048x2048xbf16> loc(#loc3571)
      %34254 = stablehlo.convert %34169 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3564)
      %34255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34256 = stablehlo.broadcast_in_dim %34255, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3565)
      %34257 = stablehlo.subtract %34256, %34254 : tensor<1x1x2048x1xf32> loc(#loc3565)
      %34258 = stablehlo.convert %34257 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3572)
      %34259 = stablehlo.broadcast_in_dim %34258, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3573)
      %34260 = stablehlo.multiply %34253, %34259 : tensor<1x1x2048x2048xbf16> loc(#loc3573)
      %34261 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34262 = stablehlo.broadcast_in_dim %34261, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3574)
      %34263 = stablehlo.multiply %34260, %34262 : tensor<1x1x2048x2048xbf16> loc(#loc3574)
      %34264 = stablehlo.add %34168, %34263 : tensor<1x1x2048x2048xbf16> loc(#loc3575)
      %34265 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3576)
      %34266 = stablehlo.convert %34264 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3577)
      %34267 = stablehlo.convert %34265 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3578)
      %34268 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34269 = stablehlo.broadcast_in_dim %34268, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3579)
      %34270 = stablehlo.subtract %34269, %34267 : tensor<1x1x2048x1xf32> loc(#loc3579)
      %34271 = stablehlo.convert %34270 : tensor<1x1x2048x1xf32> loc(#loc3580)
      %34272 = stablehlo.broadcast_in_dim %34271, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34273 = stablehlo.multiply %34266, %34272 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34274 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3582)
      %34275 = stablehlo.reduce(%34273 init: %34274) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3582)
        stablehlo.return %39096 : tensor<f32> loc(#loc3582)
      } loc(#loc3582)
      %34276 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34277 = stablehlo.broadcast_in_dim %34276, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3583)
      %34278 = stablehlo.convert %34270 : tensor<1x1x2048x1xf32> loc(#loc3580)
      %34279 = stablehlo.broadcast_in_dim %34277, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3584)
      %34280 = stablehlo.broadcast_in_dim %34278, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34281 = stablehlo.multiply %34279, %34280 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34282 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3582)
      %34283 = stablehlo.reduce(%34281 init: %34282) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3582)
        stablehlo.return %39096 : tensor<f32> loc(#loc3582)
      } loc(#loc3582)
      %34284 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34285 = stablehlo.broadcast_in_dim %34284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3585)
      %34286 = stablehlo.maximum %34285, %34283 : tensor<1xf32> loc(#loc3585)
      %34287 = stablehlo.divide %34275, %34286 : tensor<1xf32> loc(#loc3586)
      %34288 = stablehlo.broadcast_in_dim %34287, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3587)
      %34289 = stablehlo.broadcast_in_dim %34288, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3588)
      %34290 = stablehlo.subtract %34266, %34289 : tensor<1x1x2048x2048xf32> loc(#loc3588)
      %34291 = stablehlo.convert %34270 : tensor<1x1x2048x1xf32> loc(#loc3580)
      %34292 = stablehlo.broadcast_in_dim %34291, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34293 = stablehlo.multiply %34290, %34292 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34294 = stablehlo.multiply %34293, %34293 : tensor<1x1x2048x2048xf32> loc(#loc3589)
      %34295 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %34296 = stablehlo.broadcast_in_dim %34295, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3589)
      %34297 = stablehlo.multiply %34296, %34293 : tensor<1x1x2048x2048xf32> loc(#loc3589)
      %34298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3582)
      %34299 = stablehlo.reduce(%34294 init: %34298) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3582)
        stablehlo.return %39096 : tensor<f32> loc(#loc3582)
      } loc(#loc3582)
      %34300 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34301 = stablehlo.broadcast_in_dim %34300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3585)
      %34302 = stablehlo.maximum %34301, %34283 : tensor<1xf32> loc(#loc3585)
      %34303 = stablehlo.divide %34299, %34302 : tensor<1xf32> loc(#loc3586)
      %34304 = stablehlo.sqrt %34303 : tensor<1xf32> loc(#loc3590)
      %34305 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %34306 = stablehlo.broadcast_in_dim %34305, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3591)
      %34307 = stablehlo.divide %34306, %34304 : tensor<1xf32> loc(#loc3591)
      %34308 = stablehlo.convert %34270 : tensor<1x1x2048x1xf32> loc(#loc3580)
      %34309 = stablehlo.broadcast_in_dim %34308, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34310 = stablehlo.multiply %34266, %34309 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %34311 = stablehlo.abs %34310 : tensor<1x1x2048x2048xf32> loc(#loc3592)
      %34312 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34313 = stablehlo.broadcast_in_dim %34312, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3593)
      %34314 = stablehlo.compare  GE, %34310, %34313,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3593)
      %34315 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3594)
      %34316 = stablehlo.reduce(%34311 init: %34315) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3594)
        stablehlo.return %39096 : tensor<f32> loc(#loc3594)
      } loc(#loc3594)
      %34317 = stablehlo.reshape %34316 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3595)
      %34318 = stablehlo.broadcast_in_dim %34317, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3596)
      %34319 = stablehlo.compare  EQ, %34311, %34318,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3596)
      %34320 = stablehlo.convert %34319 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3597)
      %34321 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3598)
      %34322 = stablehlo.reduce(%34320 init: %34321) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3598)
        stablehlo.return %39096 : tensor<f32> loc(#loc3598)
      } loc(#loc3598)
      %34323 = stablehlo.convert %34264 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3599)
      %34324 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3600)
      %34325 = stablehlo.reduce(%34323 init: %34324) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3600)
        stablehlo.return %39096 : tensor<f32> loc(#loc3600)
      } loc(#loc3600)
      %34326 = stablehlo.broadcast_in_dim %34325, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3601)
      %34327 = stablehlo.convert %34326 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3602)
      %34328 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34329 = stablehlo.broadcast_in_dim %34328, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3603)
      %34330 = stablehlo.divide %34327, %34329 : tensor<1x1x2048x1xbf16> loc(#loc3603)
      %34331 = stablehlo.broadcast_in_dim %34330, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3604)
      %34332 = stablehlo.subtract %34264, %34331 : tensor<1x1x2048x2048xbf16> loc(#loc3604)
      %34333 = stablehlo.multiply %34332, %34332 : tensor<1x1x2048x2048xbf16> loc(#loc3605)
      %34334 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %34335 = stablehlo.broadcast_in_dim %34334, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3605)
      %34336 = stablehlo.multiply %34335, %34332 : tensor<1x1x2048x2048xbf16> loc(#loc3605)
      %34337 = stablehlo.convert %34333 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3599)
      %34338 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3600)
      %34339 = stablehlo.reduce(%34337 init: %34338) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3600)
        stablehlo.return %39096 : tensor<f32> loc(#loc3600)
      } loc(#loc3600)
      %34340 = stablehlo.broadcast_in_dim %34339, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3601)
      %34341 = stablehlo.convert %34340 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3602)
      %34342 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34343 = stablehlo.broadcast_in_dim %34342, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3603)
      %34344 = stablehlo.divide %34341, %34343 : tensor<1x1x2048x1xbf16> loc(#loc3603)
      %34345 = stablehlo.broadcast_in_dim %34330, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3604)
      %34346 = stablehlo.subtract %34264, %34345 : tensor<1x1x2048x2048xbf16> loc(#loc3604)
      %34347 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %34348 = stablehlo.broadcast_in_dim %34347, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3606)
      %34349 = stablehlo.add %34344, %34348 : tensor<1x1x2048x1xbf16> loc(#loc3606)
      %34350 = stablehlo.rsqrt %34349 : tensor<1x1x2048x1xbf16> loc(#loc3607)
      %34351 = stablehlo.divide %34350, %34349 : tensor<1x1x2048x1xbf16> loc(#loc3608)
      %34352 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34353 = stablehlo.broadcast_in_dim %34352, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3609)
      %34354 = stablehlo.multiply %34353, %34351 : tensor<1x1x2048x1xbf16> loc(#loc3609)
      %34355 = stablehlo.broadcast_in_dim %34350, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3609)
      %34356 = stablehlo.multiply %34346, %34355 : tensor<1x1x2048x2048xbf16> loc(#loc3609)
      %34357 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34358 = stablehlo.broadcast_in_dim %34357, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3610)
      %34359 = stablehlo.add %34358, %30572#0 : tensor<1x2048xbf16> loc(#loc3610)
      %34360 = stablehlo.broadcast_in_dim %34359, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3611)
      %34361 = stablehlo.broadcast_in_dim %34360, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3612)
      %34362 = stablehlo.multiply %34356, %34361 : tensor<1x1x2048x2048xbf16> loc(#loc3612)
      %34363 = stablehlo.broadcast_in_dim %30571#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3613)
      %34364 = stablehlo.broadcast_in_dim %34363, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3614)
      %34365 = stablehlo.add %34362, %34364 : tensor<1x1x2048x2048xbf16> loc(#loc3614)
      %34366 = func.call @_einsum_766(%34365, %30573#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3615)
      %34367 = "stablehlo.slice"(%34366) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3616)
      %34368 = stablehlo.reshape %34367 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3617)
      %34369 = "stablehlo.slice"(%34366) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3618)
      %34370 = stablehlo.reshape %34369 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3617)
      %34371 = "stablehlo.slice"(%34366) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3619)
      %34372 = stablehlo.reshape %34371 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3617)
      %34373 = arith.select %ptrue, %34368, %34368 : tensor<1x1x2048x32x64xbf16>
      %34374 = arith.select %ptrue, %34370, %34370 : tensor<1x1x2048x32x64xbf16>
      %34375 = arith.select %ptrue, %34372, %34372 : tensor<1x1x2048x32x64xbf16>
      %34376:2 = func.call @softplus_767(%30574#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3621)
      %34377 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %34378 = stablehlo.broadcast_in_dim %34377, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3622)
      %34379 = stablehlo.multiply %34378, %34376#0 : tensor<1x64xbf16> loc(#loc3622)
      %34380 = stablehlo.broadcast_in_dim %34379, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3623)
      %34381 = stablehlo.broadcast_in_dim %34380, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3624)
      %34382 = stablehlo.multiply %34373, %34381 : tensor<1x1x2048x32x64xbf16> loc(#loc3624)
      %34383 = func.call @_einsum_769(%34382, %34374) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3625)
      %34384 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %34385 = stablehlo.broadcast_in_dim %34384, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3626)
      %34386 = stablehlo.divide %34383, %34385 : tensor<1x1x32x2048x2048xbf16> loc(#loc3626)
      %34387 = stablehlo.tanh %34386 : tensor<1x1x32x2048x2048xbf16> loc(#loc3627)
      %34388 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34389 = stablehlo.broadcast_in_dim %34388, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3628)
      %34390 = stablehlo.subtract %34389, %34387 : tensor<1x1x32x2048x2048xbf16> loc(#loc3628)
      %34391 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %34392 = stablehlo.broadcast_in_dim %34391, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3629)
      %34393 = stablehlo.multiply %34392, %34387 : tensor<1x1x32x2048x2048xbf16> loc(#loc3629)
      %34394 = stablehlo.convert %34393 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3630)
      %34395 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %34396 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %34397 = stablehlo.multiply %34395, %34396 : tensor<f32> loc(#loc3631)
      %34398 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3632)
      %34399 = stablehlo.broadcast_in_dim %34397, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3633)
      %34400 = stablehlo.compare  GE, %34398, %34399,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3633)
      %34401 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %34402:2 = func.call @_where_770(%34400, %34394, %34401) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3634)
      %34403 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3635)
      %34404 = stablehlo.reduce(%34402#0 init: %34403) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3635)
        stablehlo.return %39096 : tensor<f32> loc(#loc3635)
      } loc(#loc3635)
      %34405 = stablehlo.broadcast_in_dim %34404, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3636)
      %34406 = stablehlo.broadcast_in_dim %34405, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3637)
      %34407 = stablehlo.subtract %34402#0, %34406 : tensor<1x1x32x2048x2048xf32> loc(#loc3637)
      %34408 = stablehlo.exponential %34407 : tensor<1x1x32x2048x2048xf32> loc(#loc3638)
      %34409 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3639)
      %34410 = stablehlo.reduce(%34408 init: %34409) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3639)
        stablehlo.return %39096 : tensor<f32> loc(#loc3639)
      } loc(#loc3639)
      %34411 = stablehlo.broadcast_in_dim %34410, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3636)
      %34412 = stablehlo.broadcast_in_dim %34411, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3640)
      %34413 = stablehlo.divide %34408, %34412 : tensor<1x1x32x2048x2048xf32> loc(#loc3640)
      %34414 = stablehlo.multiply %34411, %34411 : tensor<1x1x32x2048x1xf32> loc(#loc3641)
      %34415 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3642)
      %34416 = stablehlo.broadcast_in_dim %34415, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3643)
      %34417 = stablehlo.divide %34416, %34414 : tensor<1x1x32x2048x1xf32> loc(#loc3640)
      %34418 = stablehlo.convert %34413 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3644)
      %34419 = func.call @_einsum_771(%34418, %34375) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3645)
      %34420 = arith.select %ptrue, %34419, %34419 : tensor<1x1x2048x32x64xbf16>
      %34421 = func.call @_einsum_772(%34420, %30575#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3646)
      %34422 = arith.select %ptrue, %34421, %34421 : tensor<1x1x2048x2048xbf16>
      %34423 = stablehlo.add %34422, %34264 : tensor<1x1x2048x2048xbf16> loc(#loc3648)
      %34424 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3649)
      %34425 = stablehlo.convert %34423 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3650)
      %34426 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3651)
      %34427 = stablehlo.reduce(%34425 init: %34426) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3651)
        stablehlo.return %39096 : tensor<f32> loc(#loc3651)
      } loc(#loc3651)
      %34428 = stablehlo.broadcast_in_dim %34427, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3652)
      %34429 = stablehlo.convert %34428 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3653)
      %34430 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34431 = stablehlo.broadcast_in_dim %34430, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3654)
      %34432 = stablehlo.divide %34429, %34431 : tensor<1x1x2048x1xbf16> loc(#loc3654)
      %34433 = stablehlo.broadcast_in_dim %34432, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3655)
      %34434 = stablehlo.subtract %34423, %34433 : tensor<1x1x2048x2048xbf16> loc(#loc3655)
      %34435 = stablehlo.multiply %34434, %34434 : tensor<1x1x2048x2048xbf16> loc(#loc3656)
      %34436 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %34437 = stablehlo.broadcast_in_dim %34436, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3656)
      %34438 = stablehlo.multiply %34437, %34434 : tensor<1x1x2048x2048xbf16> loc(#loc3656)
      %34439 = stablehlo.convert %34435 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3650)
      %34440 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3651)
      %34441 = stablehlo.reduce(%34439 init: %34440) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3651)
        stablehlo.return %39096 : tensor<f32> loc(#loc3651)
      } loc(#loc3651)
      %34442 = stablehlo.broadcast_in_dim %34441, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3652)
      %34443 = stablehlo.convert %34442 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3653)
      %34444 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34445 = stablehlo.broadcast_in_dim %34444, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3654)
      %34446 = stablehlo.divide %34443, %34445 : tensor<1x1x2048x1xbf16> loc(#loc3654)
      %34447 = stablehlo.broadcast_in_dim %34432, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3655)
      %34448 = stablehlo.subtract %34423, %34447 : tensor<1x1x2048x2048xbf16> loc(#loc3655)
      %34449 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %34450 = stablehlo.broadcast_in_dim %34449, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3657)
      %34451 = stablehlo.add %34446, %34450 : tensor<1x1x2048x1xbf16> loc(#loc3657)
      %34452 = stablehlo.rsqrt %34451 : tensor<1x1x2048x1xbf16> loc(#loc3658)
      %34453 = stablehlo.divide %34452, %34451 : tensor<1x1x2048x1xbf16> loc(#loc3659)
      %34454 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34455 = stablehlo.broadcast_in_dim %34454, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3660)
      %34456 = stablehlo.multiply %34455, %34453 : tensor<1x1x2048x1xbf16> loc(#loc3660)
      %34457 = stablehlo.broadcast_in_dim %34452, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3660)
      %34458 = stablehlo.multiply %34448, %34457 : tensor<1x1x2048x2048xbf16> loc(#loc3660)
      %34459 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34460 = stablehlo.broadcast_in_dim %34459, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3661)
      %34461 = stablehlo.add %34460, %30570#0 : tensor<1x2048xbf16> loc(#loc3661)
      %34462 = stablehlo.broadcast_in_dim %34461, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3662)
      %34463 = stablehlo.broadcast_in_dim %34462, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3663)
      %34464 = stablehlo.multiply %34458, %34463 : tensor<1x1x2048x2048xbf16> loc(#loc3663)
      %34465 = stablehlo.broadcast_in_dim %30569#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3664)
      %34466 = stablehlo.broadcast_in_dim %34465, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3665)
      %34467 = stablehlo.add %34464, %34466 : tensor<1x1x2048x2048xbf16> loc(#loc3665)
      %34468 = func.call @_einsum_773(%34467, %30566#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3666)
      %34469 = arith.select %ptrue, %34468, %34468 : tensor<1x1x2048x8192xbf16>
      %34470 = stablehlo.broadcast_in_dim %30565#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3668)
      %34471 = stablehlo.broadcast_in_dim %34470, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3669)
      %34472 = stablehlo.add %34469, %34471 : tensor<1x1x2048x8192xbf16> loc(#loc3669)
      %34473 = stablehlo.multiply %34472, %34472 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34474 = stablehlo.multiply %34472, %34473 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34475 = stablehlo.multiply %34472, %34472 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34476 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %34477 = stablehlo.broadcast_in_dim %34476, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34478 = stablehlo.multiply %34477, %34475 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34479 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %34480 = stablehlo.broadcast_in_dim %34479, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34481 = stablehlo.multiply %34480, %34474 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34482 = stablehlo.add %34472, %34481 : tensor<1x1x2048x8192xbf16> loc(#loc3671)
      %34483 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %34484 = stablehlo.broadcast_in_dim %34483, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34485 = stablehlo.multiply %34484, %34482 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34486 = stablehlo.tanh %34485 : tensor<1x1x2048x8192xbf16> loc(#loc3672)
      %34487 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34488 = stablehlo.broadcast_in_dim %34487, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3673)
      %34489 = stablehlo.subtract %34488, %34486 : tensor<1x1x2048x8192xbf16> loc(#loc3673)
      %34490 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34491 = stablehlo.broadcast_in_dim %34490, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3671)
      %34492 = stablehlo.add %34491, %34486 : tensor<1x1x2048x8192xbf16> loc(#loc3671)
      %34493 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34494 = stablehlo.broadcast_in_dim %34493, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34495 = stablehlo.multiply %34494, %34492 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34496 = stablehlo.multiply %34472, %34495 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %34497 = stablehlo.convert %34424 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3674)
      %34498 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34499 = stablehlo.broadcast_in_dim %34498, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3675)
      %34500 = stablehlo.subtract %34499, %34497 : tensor<1x1x2048x1xf32> loc(#loc3675)
      %34501 = stablehlo.convert %34500 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3676)
      %34502 = stablehlo.broadcast_in_dim %34501, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3677)
      %34503 = stablehlo.multiply %34496, %34502 : tensor<1x1x2048x8192xbf16> loc(#loc3677)
      %34504 = func.call @_einsum_774(%34503, %30568#0) : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3678)
      %34505 = arith.select %ptrue, %34504, %34504 : tensor<1x1x2048x2048xbf16>
      %34506 = stablehlo.broadcast_in_dim %30567#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3680)
      %34507 = stablehlo.broadcast_in_dim %34506, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3681)
      %34508 = stablehlo.add %34505, %34507 : tensor<1x1x2048x2048xbf16> loc(#loc3681)
      %34509 = stablehlo.convert %34424 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3674)
      %34510 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34511 = stablehlo.broadcast_in_dim %34510, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3675)
      %34512 = stablehlo.subtract %34511, %34509 : tensor<1x1x2048x1xf32> loc(#loc3675)
      %34513 = stablehlo.convert %34512 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3682)
      %34514 = stablehlo.broadcast_in_dim %34513, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3683)
      %34515 = stablehlo.multiply %34508, %34514 : tensor<1x1x2048x2048xbf16> loc(#loc3683)
      %34516 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34517 = stablehlo.broadcast_in_dim %34516, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3684)
      %34518 = stablehlo.multiply %34515, %34517 : tensor<1x1x2048x2048xbf16> loc(#loc3684)
      %34519 = stablehlo.add %34423, %34518 : tensor<1x1x2048x2048xbf16> loc(#loc3685)
      %34520 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3686)
      %34521 = stablehlo.convert %34519 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3687)
      %34522 = stablehlo.convert %34520 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3688)
      %34523 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34524 = stablehlo.broadcast_in_dim %34523, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3689)
      %34525 = stablehlo.subtract %34524, %34522 : tensor<1x1x2048x1xf32> loc(#loc3689)
      %34526 = stablehlo.convert %34525 : tensor<1x1x2048x1xf32> loc(#loc3690)
      %34527 = stablehlo.broadcast_in_dim %34526, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34528 = stablehlo.multiply %34521, %34527 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34529 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3692)
      %34530 = stablehlo.reduce(%34528 init: %34529) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3692)
        stablehlo.return %39096 : tensor<f32> loc(#loc3692)
      } loc(#loc3692)
      %34531 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34532 = stablehlo.broadcast_in_dim %34531, dims = [] : (tensor<f32>) -> tensor<1x2048x2048xf32> loc(#loc3693)
      %34533 = stablehlo.convert %34525 : tensor<1x1x2048x1xf32> loc(#loc3690)
      %34534 = stablehlo.broadcast_in_dim %34532, dims = [1, 2, 3] : (tensor<1x2048x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3694)
      %34535 = stablehlo.broadcast_in_dim %34533, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34536 = stablehlo.multiply %34534, %34535 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34537 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3692)
      %34538 = stablehlo.reduce(%34536 init: %34537) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3692)
        stablehlo.return %39096 : tensor<f32> loc(#loc3692)
      } loc(#loc3692)
      %34539 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34540 = stablehlo.broadcast_in_dim %34539, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3695)
      %34541 = stablehlo.maximum %34540, %34538 : tensor<1xf32> loc(#loc3695)
      %34542 = stablehlo.divide %34530, %34541 : tensor<1xf32> loc(#loc3696)
      %34543 = stablehlo.broadcast_in_dim %34542, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3697)
      %34544 = stablehlo.broadcast_in_dim %34543, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3698)
      %34545 = stablehlo.subtract %34521, %34544 : tensor<1x1x2048x2048xf32> loc(#loc3698)
      %34546 = stablehlo.convert %34525 : tensor<1x1x2048x1xf32> loc(#loc3690)
      %34547 = stablehlo.broadcast_in_dim %34546, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34548 = stablehlo.multiply %34545, %34547 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34549 = stablehlo.multiply %34548, %34548 : tensor<1x1x2048x2048xf32> loc(#loc3699)
      %34550 = stablehlo.constant dense<2.000000e+00> : tensor<f32> loc(#loc1984)
      %34551 = stablehlo.broadcast_in_dim %34550, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3699)
      %34552 = stablehlo.multiply %34551, %34548 : tensor<1x1x2048x2048xf32> loc(#loc3699)
      %34553 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3692)
      %34554 = stablehlo.reduce(%34549 init: %34553) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3692)
        stablehlo.return %39096 : tensor<f32> loc(#loc3692)
      } loc(#loc3692)
      %34555 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34556 = stablehlo.broadcast_in_dim %34555, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3695)
      %34557 = stablehlo.maximum %34556, %34538 : tensor<1xf32> loc(#loc3695)
      %34558 = stablehlo.divide %34554, %34557 : tensor<1xf32> loc(#loc3696)
      %34559 = stablehlo.sqrt %34558 : tensor<1xf32> loc(#loc3700)
      %34560 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %34561 = stablehlo.broadcast_in_dim %34560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3701)
      %34562 = stablehlo.divide %34561, %34559 : tensor<1xf32> loc(#loc3701)
      %34563 = stablehlo.convert %34525 : tensor<1x1x2048x1xf32> loc(#loc3690)
      %34564 = stablehlo.broadcast_in_dim %34563, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34565 = stablehlo.multiply %34521, %34564 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %34566 = stablehlo.abs %34565 : tensor<1x1x2048x2048xf32> loc(#loc3702)
      %34567 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34568 = stablehlo.broadcast_in_dim %34567, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3703)
      %34569 = stablehlo.compare  GE, %34565, %34568,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3703)
      %34570 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3704)
      %34571 = stablehlo.reduce(%34566 init: %34570) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3704)
        stablehlo.return %39096 : tensor<f32> loc(#loc3704)
      } loc(#loc3704)
      %34572 = stablehlo.reshape %34571 : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3705)
      %34573 = stablehlo.broadcast_in_dim %34572, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3706)
      %34574 = stablehlo.compare  EQ, %34566, %34573,  FLOAT : (tensor<1x1x2048x2048xf32>, tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xi1> loc(#loc3706)
      %34575 = stablehlo.convert %34574 : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xf32> loc(#loc3707)
      %34576 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3708)
      %34577 = stablehlo.reduce(%34575 init: %34576) across dimensions = [1, 2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3708)
        stablehlo.return %39096 : tensor<f32> loc(#loc3708)
      } loc(#loc3708)
      %34578 = stablehlo.convert %34519 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3709)
      %34579 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3710)
      %34580 = stablehlo.reduce(%34578 init: %34579) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3710)
        stablehlo.return %39096 : tensor<f32> loc(#loc3710)
      } loc(#loc3710)
      %34581 = stablehlo.broadcast_in_dim %34580, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3711)
      %34582 = stablehlo.convert %34581 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3712)
      %34583 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34584 = stablehlo.broadcast_in_dim %34583, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3713)
      %34585 = stablehlo.divide %34582, %34584 : tensor<1x1x2048x1xbf16> loc(#loc3713)
      %34586 = stablehlo.broadcast_in_dim %34585, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3714)
      %34587 = stablehlo.subtract %34519, %34586 : tensor<1x1x2048x2048xbf16> loc(#loc3714)
      %34588 = stablehlo.multiply %34587, %34587 : tensor<1x1x2048x2048xbf16> loc(#loc3715)
      %34589 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %34590 = stablehlo.broadcast_in_dim %34589, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3715)
      %34591 = stablehlo.multiply %34590, %34587 : tensor<1x1x2048x2048xbf16> loc(#loc3715)
      %34592 = stablehlo.convert %34588 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3709)
      %34593 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3710)
      %34594 = stablehlo.reduce(%34592 init: %34593) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3710)
        stablehlo.return %39096 : tensor<f32> loc(#loc3710)
      } loc(#loc3710)
      %34595 = stablehlo.broadcast_in_dim %34594, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3711)
      %34596 = stablehlo.convert %34595 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3712)
      %34597 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34598 = stablehlo.broadcast_in_dim %34597, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3713)
      %34599 = stablehlo.divide %34596, %34598 : tensor<1x1x2048x1xbf16> loc(#loc3713)
      %34600 = stablehlo.broadcast_in_dim %34585, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3714)
      %34601 = stablehlo.subtract %34519, %34600 : tensor<1x1x2048x2048xbf16> loc(#loc3714)
      %34602 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %34603 = stablehlo.broadcast_in_dim %34602, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3716)
      %34604 = stablehlo.add %34599, %34603 : tensor<1x1x2048x1xbf16> loc(#loc3716)
      %34605 = stablehlo.rsqrt %34604 : tensor<1x1x2048x1xbf16> loc(#loc3717)
      %34606 = stablehlo.divide %34605, %34604 : tensor<1x1x2048x1xbf16> loc(#loc3718)
      %34607 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34608 = stablehlo.broadcast_in_dim %34607, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3719)
      %34609 = stablehlo.multiply %34608, %34606 : tensor<1x1x2048x1xbf16> loc(#loc3719)
      %34610 = stablehlo.broadcast_in_dim %34605, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3719)
      %34611 = stablehlo.multiply %34601, %34610 : tensor<1x1x2048x2048xbf16> loc(#loc3719)
      %34612 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34613 = stablehlo.broadcast_in_dim %34612, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3720)
      %34614 = stablehlo.add %34613, %30583#0 : tensor<1x2048xbf16> loc(#loc3720)
      %34615 = stablehlo.broadcast_in_dim %34614, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3721)
      %34616 = stablehlo.broadcast_in_dim %34615, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3722)
      %34617 = stablehlo.multiply %34611, %34616 : tensor<1x1x2048x2048xbf16> loc(#loc3722)
      %34618 = stablehlo.broadcast_in_dim %30582#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3723)
      %34619 = stablehlo.broadcast_in_dim %34618, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3724)
      %34620 = stablehlo.add %34617, %34619 : tensor<1x1x2048x2048xbf16> loc(#loc3724)
      %34621 = func.call @_einsum_775(%34620, %30584#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc3725)
      %34622 = "stablehlo.slice"(%34621) {limit_indices = dense<[1, 1, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<0> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3726)
      %34623 = stablehlo.reshape %34622 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3727)
      %34624 = "stablehlo.slice"(%34621) {limit_indices = dense<[1, 2, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 1, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3728)
      %34625 = stablehlo.reshape %34624 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3727)
      %34626 = "stablehlo.slice"(%34621) {limit_indices = dense<[1, 3, 1, 2048, 32, 64]> : tensor<6xi64>, start_indices = dense<[0, 2, 0, 0, 0, 0]> : tensor<6xi64>, strides = dense<1> : tensor<6xi64>} : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc3729)
      %34627 = stablehlo.reshape %34626 : (tensor<1x1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3727)
      %34628 = arith.select %ptrue, %34623, %34623 : tensor<1x1x2048x32x64xbf16>
      %34629 = arith.select %ptrue, %34625, %34625 : tensor<1x1x2048x32x64xbf16>
      %34630 = arith.select %ptrue, %34627, %34627 : tensor<1x1x2048x32x64xbf16>
      %34631:2 = func.call @softplus_776(%30585#0) : (tensor<1x64xbf16>) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) loc(#loc3731)
      %34632 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %34633 = stablehlo.broadcast_in_dim %34632, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3732)
      %34634 = stablehlo.multiply %34633, %34631#0 : tensor<1x64xbf16> loc(#loc3732)
      %34635 = stablehlo.broadcast_in_dim %34634, dims = [0, 4] : (tensor<1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc3733)
      %34636 = stablehlo.broadcast_in_dim %34635, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3734)
      %34637 = stablehlo.multiply %34628, %34636 : tensor<1x1x2048x32x64xbf16> loc(#loc3734)
      %34638 = func.call @_einsum_778(%34637, %34629) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3735)
      %34639 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %34640 = stablehlo.broadcast_in_dim %34639, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3736)
      %34641 = stablehlo.divide %34638, %34640 : tensor<1x1x32x2048x2048xbf16> loc(#loc3736)
      %34642 = stablehlo.tanh %34641 : tensor<1x1x32x2048x2048xbf16> loc(#loc3737)
      %34643 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34644 = stablehlo.broadcast_in_dim %34643, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3738)
      %34645 = stablehlo.subtract %34644, %34642 : tensor<1x1x32x2048x2048xbf16> loc(#loc3738)
      %34646 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %34647 = stablehlo.broadcast_in_dim %34646, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3739)
      %34648 = stablehlo.multiply %34647, %34642 : tensor<1x1x32x2048x2048xbf16> loc(#loc3739)
      %34649 = stablehlo.convert %34648 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3740)
      %34650 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %34651 = stablehlo.constant dense<5.000000e-01> : tensor<f32> loc(#loc1984)
      %34652 = stablehlo.multiply %34650, %34651 : tensor<f32> loc(#loc3741)
      %34653 = stablehlo.convert %30694 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3742)
      %34654 = stablehlo.broadcast_in_dim %34652, dims = [] : (tensor<f32>) -> tensor<1x1x1x2048x2048xf32> loc(#loc3743)
      %34655 = stablehlo.compare  GE, %34653, %34654,  FLOAT : (tensor<1x1x1x2048x2048xf32>, tensor<1x1x1x2048x2048xf32>) -> tensor<1x1x1x2048x2048xi1> loc(#loc3743)
      %34656 = stablehlo.constant dense<-2.38197633E+38> : tensor<f32> loc(#loc1984)
      %34657:2 = func.call @_where_779(%34655, %34649, %34656) : (tensor<1x1x1x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) loc(#loc3744)
      %34658 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc3745)
      %34659 = stablehlo.reduce(%34657#0 init: %34658) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc3745)
        stablehlo.return %39096 : tensor<f32> loc(#loc3745)
      } loc(#loc3745)
      %34660 = stablehlo.broadcast_in_dim %34659, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3746)
      %34661 = stablehlo.broadcast_in_dim %34660, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3747)
      %34662 = stablehlo.subtract %34657#0, %34661 : tensor<1x1x32x2048x2048xf32> loc(#loc3747)
      %34663 = stablehlo.exponential %34662 : tensor<1x1x32x2048x2048xf32> loc(#loc3748)
      %34664 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3749)
      %34665 = stablehlo.reduce(%34663 init: %34664) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3749)
        stablehlo.return %39096 : tensor<f32> loc(#loc3749)
      } loc(#loc3749)
      %34666 = stablehlo.broadcast_in_dim %34665, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3746)
      %34667 = stablehlo.broadcast_in_dim %34666, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3750)
      %34668 = stablehlo.divide %34663, %34667 : tensor<1x1x32x2048x2048xf32> loc(#loc3750)
      %34669 = stablehlo.multiply %34666, %34666 : tensor<1x1x32x2048x1xf32> loc(#loc3751)
      %34670 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc3752)
      %34671 = stablehlo.broadcast_in_dim %34670, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x1xf32> loc(#loc3753)
      %34672 = stablehlo.divide %34671, %34669 : tensor<1x1x32x2048x1xf32> loc(#loc3750)
      %34673 = stablehlo.convert %34668 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3754)
      %34674 = func.call @_einsum_780(%34673, %34630) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3755)
      %34675 = arith.select %ptrue, %34674, %34674 : tensor<1x1x2048x32x64xbf16>
      %34676 = func.call @_einsum_781(%34675, %30586#0) : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3756)
      %34677 = arith.select %ptrue, %34676, %34676 : tensor<1x1x2048x2048xbf16>
      %34678 = stablehlo.add %34677, %34519 : tensor<1x1x2048x2048xbf16> loc(#loc3758)
      %34679 = stablehlo.broadcast_in_dim %30497, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x1x2048x1xi32> loc(#loc3759)
      %34680 = stablehlo.convert %34678 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3760)
      %34681 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3761)
      %34682 = stablehlo.reduce(%34680 init: %34681) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3761)
        stablehlo.return %39096 : tensor<f32> loc(#loc3761)
      } loc(#loc3761)
      %34683 = stablehlo.broadcast_in_dim %34682, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3762)
      %34684 = stablehlo.convert %34683 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3763)
      %34685 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34686 = stablehlo.broadcast_in_dim %34685, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3764)
      %34687 = stablehlo.divide %34684, %34686 : tensor<1x1x2048x1xbf16> loc(#loc3764)
      %34688 = stablehlo.broadcast_in_dim %34687, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3765)
      %34689 = stablehlo.subtract %34678, %34688 : tensor<1x1x2048x2048xbf16> loc(#loc3765)
      %34690 = stablehlo.multiply %34689, %34689 : tensor<1x1x2048x2048xbf16> loc(#loc3766)
      %34691 = stablehlo.constant dense<2.000000e+00> : tensor<bf16> loc(#loc1984)
      %34692 = stablehlo.broadcast_in_dim %34691, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3766)
      %34693 = stablehlo.multiply %34692, %34689 : tensor<1x1x2048x2048xbf16> loc(#loc3766)
      %34694 = stablehlo.convert %34690 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xf32> loc(#loc3760)
      %34695 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3761)
      %34696 = stablehlo.reduce(%34694 init: %34695) across dimensions = [3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3761)
        stablehlo.return %39096 : tensor<f32> loc(#loc3761)
      } loc(#loc3761)
      %34697 = stablehlo.broadcast_in_dim %34696, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x1xf32> loc(#loc3762)
      %34698 = stablehlo.convert %34697 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3763)
      %34699 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %34700 = stablehlo.broadcast_in_dim %34699, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3764)
      %34701 = stablehlo.divide %34698, %34700 : tensor<1x1x2048x1xbf16> loc(#loc3764)
      %34702 = stablehlo.broadcast_in_dim %34687, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3765)
      %34703 = stablehlo.subtract %34678, %34702 : tensor<1x1x2048x2048xbf16> loc(#loc3765)
      %34704 = stablehlo.constant dense<9.983770e-07> : tensor<bf16> loc(#loc1984)
      %34705 = stablehlo.broadcast_in_dim %34704, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3767)
      %34706 = stablehlo.add %34701, %34705 : tensor<1x1x2048x1xbf16> loc(#loc3767)
      %34707 = stablehlo.rsqrt %34706 : tensor<1x1x2048x1xbf16> loc(#loc3768)
      %34708 = stablehlo.divide %34707, %34706 : tensor<1x1x2048x1xbf16> loc(#loc3769)
      %34709 = stablehlo.constant dense<-5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34710 = stablehlo.broadcast_in_dim %34709, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3770)
      %34711 = stablehlo.multiply %34710, %34708 : tensor<1x1x2048x1xbf16> loc(#loc3770)
      %34712 = stablehlo.broadcast_in_dim %34707, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3770)
      %34713 = stablehlo.multiply %34703, %34712 : tensor<1x1x2048x2048xbf16> loc(#loc3770)
      %34714 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34715 = stablehlo.broadcast_in_dim %34714, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc3771)
      %34716 = stablehlo.add %34715, %30581#0 : tensor<1x2048xbf16> loc(#loc3771)
      %34717 = stablehlo.broadcast_in_dim %34716, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3772)
      %34718 = stablehlo.broadcast_in_dim %34717, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3773)
      %34719 = stablehlo.multiply %34713, %34718 : tensor<1x1x2048x2048xbf16> loc(#loc3773)
      %34720 = stablehlo.broadcast_in_dim %30580#0, dims = [0, 3] : (tensor<1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc3774)
      %34721 = stablehlo.broadcast_in_dim %34720, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3775)
      %34722 = stablehlo.add %34719, %34721 : tensor<1x1x2048x2048xbf16> loc(#loc3775)
      %34723 = func.call @_einsum_782(%34722, %30577#0) : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3776)
      %34724 = arith.select %ptrue, %34723, %34723 : tensor<1x1x2048x8192xbf16>
      %34725 = stablehlo.broadcast_in_dim %30576#0, dims = [0, 3] : (tensor<1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc3778)
      %34726 = stablehlo.broadcast_in_dim %34725, dims = [0, 1, 2, 3] : (tensor<1x1x1x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3779)
      %34727 = stablehlo.add %34724, %34726 : tensor<1x1x2048x8192xbf16> loc(#loc3779)
      %34728 = stablehlo.multiply %34727, %34727 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34729 = stablehlo.multiply %34727, %34728 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34730 = stablehlo.multiply %34727, %34727 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34731 = stablehlo.constant dense<3.000000e+00> : tensor<bf16> loc(#loc1984)
      %34732 = stablehlo.broadcast_in_dim %34731, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34733 = stablehlo.multiply %34732, %34730 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34734 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %34735 = stablehlo.broadcast_in_dim %34734, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34736 = stablehlo.multiply %34735, %34729 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34737 = stablehlo.add %34727, %34736 : tensor<1x1x2048x8192xbf16> loc(#loc3781)
      %34738 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %34739 = stablehlo.broadcast_in_dim %34738, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34740 = stablehlo.multiply %34739, %34737 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34741 = stablehlo.tanh %34740 : tensor<1x1x2048x8192xbf16> loc(#loc3782)
      %34742 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34743 = stablehlo.broadcast_in_dim %34742, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3783)
      %34744 = stablehlo.subtract %34743, %34741 : tensor<1x1x2048x8192xbf16> loc(#loc3783)
      %34745 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %34746 = stablehlo.broadcast_in_dim %34745, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3781)
      %34747 = stablehlo.add %34746, %34741 : tensor<1x1x2048x8192xbf16> loc(#loc3781)
      %34748 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %34749 = stablehlo.broadcast_in_dim %34748, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34750 = stablehlo.multiply %34749, %34747 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34751 = stablehlo.multiply %34727, %34750 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %34752 = stablehlo.convert %34679 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3784)
      %34753 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34754 = stablehlo.broadcast_in_dim %34753, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3785)
      %34755 = stablehlo.subtract %34754, %34752 : tensor<1x1x2048x1xf32> loc(#loc3785)
      %34756 = stablehlo.convert %34755 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3786)
      %34757 = stablehlo.broadcast_in_dim %34756, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3787)
      %34758 = stablehlo.multiply %34751, %34757 : tensor<1x1x2048x8192xbf16> loc(#loc3787)
      %34759 = stablehlo.convert %34679 : (tensor<1x1x2048x1xi32>) -> tensor<1x1x2048x1xf32> loc(#loc3784)
      %34760 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc1984)
      %34761 = stablehlo.broadcast_in_dim %34760, dims = [] : (tensor<f32>) -> tensor<1x1x2048x1xf32> loc(#loc3785)
      %34762 = stablehlo.subtract %34761, %34759 : tensor<1x1x2048x1xf32> loc(#loc3785)
      %34763 = stablehlo.convert %34762 : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x1xbf16> loc(#loc3788)
      %34764 = func.call @_where_783(%30495, %30376) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34765 = stablehlo.multiply %34764, %33032 : tensor<1xf32> loc(#loc3790)
      %34766 = stablehlo.divide %34765, %33027 : tensor<1xf32> loc(#loc3036)
      %34767 = stablehlo.broadcast_in_dim %34766, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3791)
      %34768 = stablehlo.multiply %34767, %33022 : tensor<1x1x2048x2048xf32> loc(#loc3039)
      %34769 = stablehlo.broadcast_in_dim %33016, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %34770 = stablehlo.multiply %34768, %34769 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %34771 = stablehlo.negate %34770 : tensor<1x1x2048x2048xf32> loc(#loc3792)
      %34772 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3793)
      %34773 = stablehlo.reduce(%34771 init: %34772) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3793)
        stablehlo.return %39096 : tensor<f32> loc(#loc3793)
      } loc(#loc3793)
      %34774 = stablehlo.reshape %34773 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3794)
      %34775 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3795)
      %34776 = stablehlo.reduce(%34774 init: %34775) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3795)
        stablehlo.return %39096 : tensor<f32> loc(#loc3795)
      } loc(#loc3795)
      %34777 = stablehlo.broadcast_in_dim %34776, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3796)
      %34778 = func.call @_where_784(%30495, %30367) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34779 = stablehlo.add %34777, %34778 : tensor<1xf32> loc(#loc3797)
      %34780 = stablehlo.divide %34779, %33011 : tensor<1xf32> loc(#loc3036)
      %34781 = stablehlo.broadcast_in_dim %34780, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3791)
      %34782 = stablehlo.broadcast_in_dim %32996, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %34783 = stablehlo.multiply %34781, %34782 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %34784 = stablehlo.add %34770, %34783 : tensor<1x1x2048x2048xf32> loc(#loc3798)
      %34785 = func.call @_where_785(%30495, %30358) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34786 = stablehlo.divide %34785, %33047 : tensor<1xf32> loc(#loc3799)
      %34787 = stablehlo.broadcast_in_dim %34786, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3800)
      %34788 = stablehlo.multiply %34787, %33045 : tensor<1x1x2048x2048xf32> loc(#loc3801)
      %34789 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34790 = stablehlo.broadcast_in_dim %34789, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3802)
      %34791 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %34792 = stablehlo.broadcast_in_dim %34791, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3046)
      %34793 = stablehlo.compare  EQ, %33039, %34792,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3046)
      %34794 = stablehlo.select %34793, %34788, %34790 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3803)
      %34795 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %34796 = stablehlo.broadcast_in_dim %34795, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3046)
      %34797 = stablehlo.compare  EQ, %33039, %34796,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3046)
      %34798 = stablehlo.select %34797, %34788, %34790 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3803)
      %34799 = stablehlo.negate %34794 : tensor<1x1x2048x2048xf32> loc(#loc3804)
      %34800 = stablehlo.add %34798, %34799 : tensor<1x1x2048x2048xf32> loc(#loc3805)
      %34801 = stablehlo.broadcast_in_dim %33033, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3031)
      %34802 = stablehlo.multiply %34800, %34801 : tensor<1x1x2048x2048xf32> loc(#loc3031)
      %34803 = stablehlo.add %34784, %34802 : tensor<1x1x2048x2048xf32> loc(#loc3798)
      %34804 = stablehlo.convert %34803 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3806)
      %34805 = func.call @_where_786(%30495, %30349) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34806 = stablehlo.multiply %34805, %32777 : tensor<1xf32> loc(#loc3807)
      %34807 = stablehlo.divide %34806, %32772 : tensor<1xf32> loc(#loc2926)
      %34808 = stablehlo.broadcast_in_dim %34807, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3808)
      %34809 = stablehlo.multiply %34808, %32767 : tensor<1x1x2048x2048xf32> loc(#loc2929)
      %34810 = stablehlo.broadcast_in_dim %32761, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %34811 = stablehlo.multiply %34809, %34810 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %34812 = stablehlo.negate %34811 : tensor<1x1x2048x2048xf32> loc(#loc3809)
      %34813 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3810)
      %34814 = stablehlo.reduce(%34812 init: %34813) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3810)
        stablehlo.return %39096 : tensor<f32> loc(#loc3810)
      } loc(#loc3810)
      %34815 = stablehlo.reshape %34814 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3811)
      %34816 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3812)
      %34817 = stablehlo.reduce(%34815 init: %34816) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3812)
        stablehlo.return %39096 : tensor<f32> loc(#loc3812)
      } loc(#loc3812)
      %34818 = stablehlo.broadcast_in_dim %34817, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3813)
      %34819 = func.call @_where_787(%30495, %30340) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34820 = stablehlo.add %34818, %34819 : tensor<1xf32> loc(#loc3797)
      %34821 = stablehlo.divide %34820, %32756 : tensor<1xf32> loc(#loc2926)
      %34822 = stablehlo.broadcast_in_dim %34821, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3808)
      %34823 = stablehlo.broadcast_in_dim %32741, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %34824 = stablehlo.multiply %34822, %34823 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %34825 = stablehlo.add %34811, %34824 : tensor<1x1x2048x2048xf32> loc(#loc3814)
      %34826 = func.call @_where_788(%30495, %30331) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34827 = stablehlo.divide %34826, %32792 : tensor<1xf32> loc(#loc3815)
      %34828 = stablehlo.broadcast_in_dim %34827, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3816)
      %34829 = stablehlo.multiply %34828, %32790 : tensor<1x1x2048x2048xf32> loc(#loc3817)
      %34830 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34831 = stablehlo.broadcast_in_dim %34830, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3818)
      %34832 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %34833 = stablehlo.broadcast_in_dim %34832, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2936)
      %34834 = stablehlo.compare  EQ, %32784, %34833,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2936)
      %34835 = stablehlo.select %34834, %34829, %34831 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3819)
      %34836 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %34837 = stablehlo.broadcast_in_dim %34836, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2936)
      %34838 = stablehlo.compare  EQ, %32784, %34837,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2936)
      %34839 = stablehlo.select %34838, %34829, %34831 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3819)
      %34840 = stablehlo.negate %34835 : tensor<1x1x2048x2048xf32> loc(#loc3820)
      %34841 = stablehlo.add %34839, %34840 : tensor<1x1x2048x2048xf32> loc(#loc3821)
      %34842 = stablehlo.broadcast_in_dim %32778, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2921)
      %34843 = stablehlo.multiply %34841, %34842 : tensor<1x1x2048x2048xf32> loc(#loc2921)
      %34844 = stablehlo.add %34825, %34843 : tensor<1x1x2048x2048xf32> loc(#loc3814)
      %34845 = stablehlo.convert %34844 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3822)
      %34846 = func.call @_where_789(%30495, %30322) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34847 = stablehlo.multiply %34846, %32522 : tensor<1xf32> loc(#loc3823)
      %34848 = stablehlo.divide %34847, %32517 : tensor<1xf32> loc(#loc2816)
      %34849 = stablehlo.broadcast_in_dim %34848, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3824)
      %34850 = stablehlo.multiply %34849, %32512 : tensor<1x1x2048x2048xf32> loc(#loc2819)
      %34851 = stablehlo.broadcast_in_dim %32506, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %34852 = stablehlo.multiply %34850, %34851 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %34853 = stablehlo.negate %34852 : tensor<1x1x2048x2048xf32> loc(#loc3825)
      %34854 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3826)
      %34855 = stablehlo.reduce(%34853 init: %34854) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3826)
        stablehlo.return %39096 : tensor<f32> loc(#loc3826)
      } loc(#loc3826)
      %34856 = stablehlo.reshape %34855 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3827)
      %34857 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3828)
      %34858 = stablehlo.reduce(%34856 init: %34857) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3828)
        stablehlo.return %39096 : tensor<f32> loc(#loc3828)
      } loc(#loc3828)
      %34859 = stablehlo.broadcast_in_dim %34858, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3829)
      %34860 = func.call @_where_790(%30495, %30313) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34861 = stablehlo.add %34859, %34860 : tensor<1xf32> loc(#loc3797)
      %34862 = stablehlo.divide %34861, %32501 : tensor<1xf32> loc(#loc2816)
      %34863 = stablehlo.broadcast_in_dim %34862, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3824)
      %34864 = stablehlo.broadcast_in_dim %32486, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %34865 = stablehlo.multiply %34863, %34864 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %34866 = stablehlo.add %34852, %34865 : tensor<1x1x2048x2048xf32> loc(#loc3830)
      %34867 = func.call @_where_791(%30495, %30304) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34868 = stablehlo.divide %34867, %32537 : tensor<1xf32> loc(#loc3831)
      %34869 = stablehlo.broadcast_in_dim %34868, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3832)
      %34870 = stablehlo.multiply %34869, %32535 : tensor<1x1x2048x2048xf32> loc(#loc3833)
      %34871 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34872 = stablehlo.broadcast_in_dim %34871, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3834)
      %34873 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %34874 = stablehlo.broadcast_in_dim %34873, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2826)
      %34875 = stablehlo.compare  EQ, %32529, %34874,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2826)
      %34876 = stablehlo.select %34875, %34870, %34872 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3835)
      %34877 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %34878 = stablehlo.broadcast_in_dim %34877, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2826)
      %34879 = stablehlo.compare  EQ, %32529, %34878,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2826)
      %34880 = stablehlo.select %34879, %34870, %34872 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3835)
      %34881 = stablehlo.negate %34876 : tensor<1x1x2048x2048xf32> loc(#loc3836)
      %34882 = stablehlo.add %34880, %34881 : tensor<1x1x2048x2048xf32> loc(#loc3837)
      %34883 = stablehlo.broadcast_in_dim %32523, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2811)
      %34884 = stablehlo.multiply %34882, %34883 : tensor<1x1x2048x2048xf32> loc(#loc2811)
      %34885 = stablehlo.add %34866, %34884 : tensor<1x1x2048x2048xf32> loc(#loc3830)
      %34886 = stablehlo.convert %34885 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3838)
      %34887 = func.call @_where_792(%30495, %30295) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34888 = stablehlo.multiply %34887, %32267 : tensor<1xf32> loc(#loc3839)
      %34889 = stablehlo.divide %34888, %32262 : tensor<1xf32> loc(#loc2706)
      %34890 = stablehlo.broadcast_in_dim %34889, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3840)
      %34891 = stablehlo.multiply %34890, %32257 : tensor<1x1x2048x2048xf32> loc(#loc2709)
      %34892 = stablehlo.broadcast_in_dim %32251, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %34893 = stablehlo.multiply %34891, %34892 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %34894 = stablehlo.negate %34893 : tensor<1x1x2048x2048xf32> loc(#loc3841)
      %34895 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3842)
      %34896 = stablehlo.reduce(%34894 init: %34895) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3842)
        stablehlo.return %39096 : tensor<f32> loc(#loc3842)
      } loc(#loc3842)
      %34897 = stablehlo.reshape %34896 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3843)
      %34898 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3844)
      %34899 = stablehlo.reduce(%34897 init: %34898) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3844)
        stablehlo.return %39096 : tensor<f32> loc(#loc3844)
      } loc(#loc3844)
      %34900 = stablehlo.broadcast_in_dim %34899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3845)
      %34901 = func.call @_where_793(%30495, %30286) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34902 = stablehlo.add %34900, %34901 : tensor<1xf32> loc(#loc3797)
      %34903 = stablehlo.divide %34902, %32246 : tensor<1xf32> loc(#loc2706)
      %34904 = stablehlo.broadcast_in_dim %34903, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3840)
      %34905 = stablehlo.broadcast_in_dim %32231, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %34906 = stablehlo.multiply %34904, %34905 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %34907 = stablehlo.add %34893, %34906 : tensor<1x1x2048x2048xf32> loc(#loc3846)
      %34908 = func.call @_where_794(%30495, %30277) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34909 = stablehlo.divide %34908, %32282 : tensor<1xf32> loc(#loc3847)
      %34910 = stablehlo.broadcast_in_dim %34909, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3848)
      %34911 = stablehlo.multiply %34910, %32280 : tensor<1x1x2048x2048xf32> loc(#loc3849)
      %34912 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34913 = stablehlo.broadcast_in_dim %34912, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3850)
      %34914 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %34915 = stablehlo.broadcast_in_dim %34914, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2716)
      %34916 = stablehlo.compare  EQ, %32274, %34915,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2716)
      %34917 = stablehlo.select %34916, %34911, %34913 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3851)
      %34918 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %34919 = stablehlo.broadcast_in_dim %34918, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2716)
      %34920 = stablehlo.compare  EQ, %32274, %34919,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2716)
      %34921 = stablehlo.select %34920, %34911, %34913 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3851)
      %34922 = stablehlo.negate %34917 : tensor<1x1x2048x2048xf32> loc(#loc3852)
      %34923 = stablehlo.add %34921, %34922 : tensor<1x1x2048x2048xf32> loc(#loc3853)
      %34924 = stablehlo.broadcast_in_dim %32268, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2701)
      %34925 = stablehlo.multiply %34923, %34924 : tensor<1x1x2048x2048xf32> loc(#loc2701)
      %34926 = stablehlo.add %34907, %34925 : tensor<1x1x2048x2048xf32> loc(#loc3846)
      %34927 = stablehlo.convert %34926 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3854)
      %34928 = func.call @_where_795(%30495, %30268) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34929 = stablehlo.multiply %34928, %32012 : tensor<1xf32> loc(#loc3855)
      %34930 = stablehlo.divide %34929, %32007 : tensor<1xf32> loc(#loc2596)
      %34931 = stablehlo.broadcast_in_dim %34930, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3856)
      %34932 = stablehlo.multiply %34931, %32002 : tensor<1x1x2048x2048xf32> loc(#loc2599)
      %34933 = stablehlo.broadcast_in_dim %31996, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %34934 = stablehlo.multiply %34932, %34933 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %34935 = stablehlo.negate %34934 : tensor<1x1x2048x2048xf32> loc(#loc3857)
      %34936 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3858)
      %34937 = stablehlo.reduce(%34935 init: %34936) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3858)
        stablehlo.return %39096 : tensor<f32> loc(#loc3858)
      } loc(#loc3858)
      %34938 = stablehlo.reshape %34937 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3859)
      %34939 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3860)
      %34940 = stablehlo.reduce(%34938 init: %34939) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3860)
        stablehlo.return %39096 : tensor<f32> loc(#loc3860)
      } loc(#loc3860)
      %34941 = stablehlo.broadcast_in_dim %34940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3861)
      %34942 = func.call @_where_796(%30495, %30259) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34943 = stablehlo.add %34941, %34942 : tensor<1xf32> loc(#loc3797)
      %34944 = stablehlo.divide %34943, %31991 : tensor<1xf32> loc(#loc2596)
      %34945 = stablehlo.broadcast_in_dim %34944, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3856)
      %34946 = stablehlo.broadcast_in_dim %31976, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %34947 = stablehlo.multiply %34945, %34946 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %34948 = stablehlo.add %34934, %34947 : tensor<1x1x2048x2048xf32> loc(#loc3862)
      %34949 = func.call @_where_797(%30495, %30250) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34950 = stablehlo.divide %34949, %32027 : tensor<1xf32> loc(#loc3863)
      %34951 = stablehlo.broadcast_in_dim %34950, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3864)
      %34952 = stablehlo.multiply %34951, %32025 : tensor<1x1x2048x2048xf32> loc(#loc3865)
      %34953 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34954 = stablehlo.broadcast_in_dim %34953, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3866)
      %34955 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %34956 = stablehlo.broadcast_in_dim %34955, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2606)
      %34957 = stablehlo.compare  EQ, %32019, %34956,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2606)
      %34958 = stablehlo.select %34957, %34952, %34954 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3867)
      %34959 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %34960 = stablehlo.broadcast_in_dim %34959, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2606)
      %34961 = stablehlo.compare  EQ, %32019, %34960,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2606)
      %34962 = stablehlo.select %34961, %34952, %34954 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3867)
      %34963 = stablehlo.negate %34958 : tensor<1x1x2048x2048xf32> loc(#loc3868)
      %34964 = stablehlo.add %34962, %34963 : tensor<1x1x2048x2048xf32> loc(#loc3869)
      %34965 = stablehlo.broadcast_in_dim %32013, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2591)
      %34966 = stablehlo.multiply %34964, %34965 : tensor<1x1x2048x2048xf32> loc(#loc2591)
      %34967 = stablehlo.add %34948, %34966 : tensor<1x1x2048x2048xf32> loc(#loc3862)
      %34968 = stablehlo.convert %34967 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3870)
      %34969 = func.call @_where_798(%30495, %30241) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34970 = stablehlo.multiply %34969, %31757 : tensor<1xf32> loc(#loc3871)
      %34971 = stablehlo.divide %34970, %31752 : tensor<1xf32> loc(#loc2486)
      %34972 = stablehlo.broadcast_in_dim %34971, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3872)
      %34973 = stablehlo.multiply %34972, %31747 : tensor<1x1x2048x2048xf32> loc(#loc2489)
      %34974 = stablehlo.broadcast_in_dim %31741, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %34975 = stablehlo.multiply %34973, %34974 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %34976 = stablehlo.negate %34975 : tensor<1x1x2048x2048xf32> loc(#loc3873)
      %34977 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3874)
      %34978 = stablehlo.reduce(%34976 init: %34977) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3874)
        stablehlo.return %39096 : tensor<f32> loc(#loc3874)
      } loc(#loc3874)
      %34979 = stablehlo.reshape %34978 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3875)
      %34980 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3876)
      %34981 = stablehlo.reduce(%34979 init: %34980) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3876)
        stablehlo.return %39096 : tensor<f32> loc(#loc3876)
      } loc(#loc3876)
      %34982 = stablehlo.broadcast_in_dim %34981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3877)
      %34983 = func.call @_where_799(%30495, %30232) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34984 = stablehlo.add %34982, %34983 : tensor<1xf32> loc(#loc3797)
      %34985 = stablehlo.divide %34984, %31736 : tensor<1xf32> loc(#loc2486)
      %34986 = stablehlo.broadcast_in_dim %34985, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3872)
      %34987 = stablehlo.broadcast_in_dim %31721, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %34988 = stablehlo.multiply %34986, %34987 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %34989 = stablehlo.add %34975, %34988 : tensor<1x1x2048x2048xf32> loc(#loc3878)
      %34990 = func.call @_where_800(%30495, %30223) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %34991 = stablehlo.divide %34990, %31772 : tensor<1xf32> loc(#loc3879)
      %34992 = stablehlo.broadcast_in_dim %34991, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3880)
      %34993 = stablehlo.multiply %34992, %31770 : tensor<1x1x2048x2048xf32> loc(#loc3881)
      %34994 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %34995 = stablehlo.broadcast_in_dim %34994, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3882)
      %34996 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %34997 = stablehlo.broadcast_in_dim %34996, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2496)
      %34998 = stablehlo.compare  EQ, %31764, %34997,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2496)
      %34999 = stablehlo.select %34998, %34993, %34995 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3883)
      %35000 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35001 = stablehlo.broadcast_in_dim %35000, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2496)
      %35002 = stablehlo.compare  EQ, %31764, %35001,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2496)
      %35003 = stablehlo.select %35002, %34993, %34995 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3883)
      %35004 = stablehlo.negate %34999 : tensor<1x1x2048x2048xf32> loc(#loc3884)
      %35005 = stablehlo.add %35003, %35004 : tensor<1x1x2048x2048xf32> loc(#loc3885)
      %35006 = stablehlo.broadcast_in_dim %31758, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2481)
      %35007 = stablehlo.multiply %35005, %35006 : tensor<1x1x2048x2048xf32> loc(#loc2481)
      %35008 = stablehlo.add %34989, %35007 : tensor<1x1x2048x2048xf32> loc(#loc3878)
      %35009 = stablehlo.convert %35008 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3886)
      %35010 = func.call @_where_801(%30495, %30214) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35011 = stablehlo.multiply %35010, %31502 : tensor<1xf32> loc(#loc3887)
      %35012 = stablehlo.divide %35011, %31497 : tensor<1xf32> loc(#loc2376)
      %35013 = stablehlo.broadcast_in_dim %35012, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3888)
      %35014 = stablehlo.multiply %35013, %31492 : tensor<1x1x2048x2048xf32> loc(#loc2379)
      %35015 = stablehlo.broadcast_in_dim %31486, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %35016 = stablehlo.multiply %35014, %35015 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %35017 = stablehlo.negate %35016 : tensor<1x1x2048x2048xf32> loc(#loc3889)
      %35018 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3890)
      %35019 = stablehlo.reduce(%35017 init: %35018) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3890)
        stablehlo.return %39096 : tensor<f32> loc(#loc3890)
      } loc(#loc3890)
      %35020 = stablehlo.reshape %35019 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3891)
      %35021 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3892)
      %35022 = stablehlo.reduce(%35020 init: %35021) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3892)
        stablehlo.return %39096 : tensor<f32> loc(#loc3892)
      } loc(#loc3892)
      %35023 = stablehlo.broadcast_in_dim %35022, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3893)
      %35024 = func.call @_where_802(%30495, %30205) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35025 = stablehlo.add %35023, %35024 : tensor<1xf32> loc(#loc3797)
      %35026 = stablehlo.divide %35025, %31481 : tensor<1xf32> loc(#loc2376)
      %35027 = stablehlo.broadcast_in_dim %35026, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3888)
      %35028 = stablehlo.broadcast_in_dim %31466, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %35029 = stablehlo.multiply %35027, %35028 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %35030 = stablehlo.add %35016, %35029 : tensor<1x1x2048x2048xf32> loc(#loc3894)
      %35031 = func.call @_where_803(%30495, %30196) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35032 = stablehlo.divide %35031, %31517 : tensor<1xf32> loc(#loc3895)
      %35033 = stablehlo.broadcast_in_dim %35032, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3896)
      %35034 = stablehlo.multiply %35033, %31515 : tensor<1x1x2048x2048xf32> loc(#loc3897)
      %35035 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35036 = stablehlo.broadcast_in_dim %35035, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3898)
      %35037 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35038 = stablehlo.broadcast_in_dim %35037, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2386)
      %35039 = stablehlo.compare  EQ, %31509, %35038,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2386)
      %35040 = stablehlo.select %35039, %35034, %35036 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3899)
      %35041 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35042 = stablehlo.broadcast_in_dim %35041, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2386)
      %35043 = stablehlo.compare  EQ, %31509, %35042,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2386)
      %35044 = stablehlo.select %35043, %35034, %35036 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3899)
      %35045 = stablehlo.negate %35040 : tensor<1x1x2048x2048xf32> loc(#loc3900)
      %35046 = stablehlo.add %35044, %35045 : tensor<1x1x2048x2048xf32> loc(#loc3901)
      %35047 = stablehlo.broadcast_in_dim %31503, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2371)
      %35048 = stablehlo.multiply %35046, %35047 : tensor<1x1x2048x2048xf32> loc(#loc2371)
      %35049 = stablehlo.add %35030, %35048 : tensor<1x1x2048x2048xf32> loc(#loc3894)
      %35050 = stablehlo.convert %35049 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3902)
      %35051 = func.call @_where_804(%30495, %30187) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35052 = stablehlo.multiply %35051, %31247 : tensor<1xf32> loc(#loc3903)
      %35053 = stablehlo.divide %35052, %31242 : tensor<1xf32> loc(#loc2266)
      %35054 = stablehlo.broadcast_in_dim %35053, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3904)
      %35055 = stablehlo.multiply %35054, %31237 : tensor<1x1x2048x2048xf32> loc(#loc2269)
      %35056 = stablehlo.broadcast_in_dim %31231, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %35057 = stablehlo.multiply %35055, %35056 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %35058 = stablehlo.negate %35057 : tensor<1x1x2048x2048xf32> loc(#loc3905)
      %35059 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3906)
      %35060 = stablehlo.reduce(%35058 init: %35059) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3906)
        stablehlo.return %39096 : tensor<f32> loc(#loc3906)
      } loc(#loc3906)
      %35061 = stablehlo.reshape %35060 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3907)
      %35062 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3908)
      %35063 = stablehlo.reduce(%35061 init: %35062) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3908)
        stablehlo.return %39096 : tensor<f32> loc(#loc3908)
      } loc(#loc3908)
      %35064 = stablehlo.broadcast_in_dim %35063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3909)
      %35065 = func.call @_where_805(%30495, %30178) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35066 = stablehlo.add %35064, %35065 : tensor<1xf32> loc(#loc3797)
      %35067 = stablehlo.divide %35066, %31226 : tensor<1xf32> loc(#loc2266)
      %35068 = stablehlo.broadcast_in_dim %35067, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3904)
      %35069 = stablehlo.broadcast_in_dim %31211, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %35070 = stablehlo.multiply %35068, %35069 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %35071 = stablehlo.add %35057, %35070 : tensor<1x1x2048x2048xf32> loc(#loc3910)
      %35072 = func.call @_where_806(%30495, %30169) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35073 = stablehlo.divide %35072, %31262 : tensor<1xf32> loc(#loc3911)
      %35074 = stablehlo.broadcast_in_dim %35073, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3912)
      %35075 = stablehlo.multiply %35074, %31260 : tensor<1x1x2048x2048xf32> loc(#loc3913)
      %35076 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35077 = stablehlo.broadcast_in_dim %35076, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3914)
      %35078 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35079 = stablehlo.broadcast_in_dim %35078, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2276)
      %35080 = stablehlo.compare  EQ, %31254, %35079,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2276)
      %35081 = stablehlo.select %35080, %35075, %35077 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3915)
      %35082 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35083 = stablehlo.broadcast_in_dim %35082, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2276)
      %35084 = stablehlo.compare  EQ, %31254, %35083,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2276)
      %35085 = stablehlo.select %35084, %35075, %35077 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3915)
      %35086 = stablehlo.negate %35081 : tensor<1x1x2048x2048xf32> loc(#loc3916)
      %35087 = stablehlo.add %35085, %35086 : tensor<1x1x2048x2048xf32> loc(#loc3917)
      %35088 = stablehlo.broadcast_in_dim %31248, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2261)
      %35089 = stablehlo.multiply %35087, %35088 : tensor<1x1x2048x2048xf32> loc(#loc2261)
      %35090 = stablehlo.add %35071, %35089 : tensor<1x1x2048x2048xf32> loc(#loc3910)
      %35091 = stablehlo.convert %35090 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3918)
      %35092 = func.call @_where_807(%30495, %30160) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35093 = stablehlo.multiply %35092, %34562 : tensor<1xf32> loc(#loc3919)
      %35094 = stablehlo.divide %35093, %34557 : tensor<1xf32> loc(#loc3696)
      %35095 = stablehlo.broadcast_in_dim %35094, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3920)
      %35096 = stablehlo.multiply %35095, %34552 : tensor<1x1x2048x2048xf32> loc(#loc3699)
      %35097 = stablehlo.broadcast_in_dim %34546, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %35098 = stablehlo.multiply %35096, %35097 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %35099 = stablehlo.negate %35098 : tensor<1x1x2048x2048xf32> loc(#loc3921)
      %35100 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3922)
      %35101 = stablehlo.reduce(%35099 init: %35100) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3922)
        stablehlo.return %39096 : tensor<f32> loc(#loc3922)
      } loc(#loc3922)
      %35102 = stablehlo.reshape %35101 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3923)
      %35103 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3924)
      %35104 = stablehlo.reduce(%35102 init: %35103) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3924)
        stablehlo.return %39096 : tensor<f32> loc(#loc3924)
      } loc(#loc3924)
      %35105 = stablehlo.broadcast_in_dim %35104, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3925)
      %35106 = func.call @_where_808(%30495, %30151) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35107 = stablehlo.add %35105, %35106 : tensor<1xf32> loc(#loc3797)
      %35108 = stablehlo.divide %35107, %34541 : tensor<1xf32> loc(#loc3696)
      %35109 = stablehlo.broadcast_in_dim %35108, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3920)
      %35110 = stablehlo.broadcast_in_dim %34526, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %35111 = stablehlo.multiply %35109, %35110 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %35112 = stablehlo.add %35098, %35111 : tensor<1x1x2048x2048xf32> loc(#loc3926)
      %35113 = func.call @_where_809(%30495, %30142) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35114 = stablehlo.divide %35113, %34577 : tensor<1xf32> loc(#loc3927)
      %35115 = stablehlo.broadcast_in_dim %35114, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3928)
      %35116 = stablehlo.multiply %35115, %34575 : tensor<1x1x2048x2048xf32> loc(#loc3929)
      %35117 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35118 = stablehlo.broadcast_in_dim %35117, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3930)
      %35119 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35120 = stablehlo.broadcast_in_dim %35119, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3706)
      %35121 = stablehlo.compare  EQ, %34569, %35120,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3706)
      %35122 = stablehlo.select %35121, %35116, %35118 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3931)
      %35123 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35124 = stablehlo.broadcast_in_dim %35123, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3706)
      %35125 = stablehlo.compare  EQ, %34569, %35124,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3706)
      %35126 = stablehlo.select %35125, %35116, %35118 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3931)
      %35127 = stablehlo.negate %35122 : tensor<1x1x2048x2048xf32> loc(#loc3932)
      %35128 = stablehlo.add %35126, %35127 : tensor<1x1x2048x2048xf32> loc(#loc3933)
      %35129 = stablehlo.broadcast_in_dim %34563, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3691)
      %35130 = stablehlo.multiply %35128, %35129 : tensor<1x1x2048x2048xf32> loc(#loc3691)
      %35131 = stablehlo.add %35112, %35130 : tensor<1x1x2048x2048xf32> loc(#loc3926)
      %35132 = stablehlo.convert %35131 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3934)
      %35133 = func.call @_where_810(%30495, %30133) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35134 = stablehlo.multiply %35133, %34307 : tensor<1xf32> loc(#loc3935)
      %35135 = stablehlo.divide %35134, %34302 : tensor<1xf32> loc(#loc3586)
      %35136 = stablehlo.broadcast_in_dim %35135, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3936)
      %35137 = stablehlo.multiply %35136, %34297 : tensor<1x1x2048x2048xf32> loc(#loc3589)
      %35138 = stablehlo.broadcast_in_dim %34291, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %35139 = stablehlo.multiply %35137, %35138 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %35140 = stablehlo.negate %35139 : tensor<1x1x2048x2048xf32> loc(#loc3937)
      %35141 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3938)
      %35142 = stablehlo.reduce(%35140 init: %35141) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3938)
        stablehlo.return %39096 : tensor<f32> loc(#loc3938)
      } loc(#loc3938)
      %35143 = stablehlo.reshape %35142 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3939)
      %35144 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3940)
      %35145 = stablehlo.reduce(%35143 init: %35144) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3940)
        stablehlo.return %39096 : tensor<f32> loc(#loc3940)
      } loc(#loc3940)
      %35146 = stablehlo.broadcast_in_dim %35145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3941)
      %35147 = func.call @_where_811(%30495, %30124) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35148 = stablehlo.add %35146, %35147 : tensor<1xf32> loc(#loc3797)
      %35149 = stablehlo.divide %35148, %34286 : tensor<1xf32> loc(#loc3586)
      %35150 = stablehlo.broadcast_in_dim %35149, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3936)
      %35151 = stablehlo.broadcast_in_dim %34271, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %35152 = stablehlo.multiply %35150, %35151 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %35153 = stablehlo.add %35139, %35152 : tensor<1x1x2048x2048xf32> loc(#loc3942)
      %35154 = func.call @_where_812(%30495, %30115) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35155 = stablehlo.divide %35154, %34322 : tensor<1xf32> loc(#loc3943)
      %35156 = stablehlo.broadcast_in_dim %35155, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3944)
      %35157 = stablehlo.multiply %35156, %34320 : tensor<1x1x2048x2048xf32> loc(#loc3945)
      %35158 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35159 = stablehlo.broadcast_in_dim %35158, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3946)
      %35160 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35161 = stablehlo.broadcast_in_dim %35160, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3596)
      %35162 = stablehlo.compare  EQ, %34314, %35161,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3596)
      %35163 = stablehlo.select %35162, %35157, %35159 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3947)
      %35164 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35165 = stablehlo.broadcast_in_dim %35164, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3596)
      %35166 = stablehlo.compare  EQ, %34314, %35165,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3596)
      %35167 = stablehlo.select %35166, %35157, %35159 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3947)
      %35168 = stablehlo.negate %35163 : tensor<1x1x2048x2048xf32> loc(#loc3948)
      %35169 = stablehlo.add %35167, %35168 : tensor<1x1x2048x2048xf32> loc(#loc3949)
      %35170 = stablehlo.broadcast_in_dim %34308, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3581)
      %35171 = stablehlo.multiply %35169, %35170 : tensor<1x1x2048x2048xf32> loc(#loc3581)
      %35172 = stablehlo.add %35153, %35171 : tensor<1x1x2048x2048xf32> loc(#loc3942)
      %35173 = stablehlo.convert %35172 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3950)
      %35174 = func.call @_where_813(%30495, %30106) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35175 = stablehlo.multiply %35174, %34052 : tensor<1xf32> loc(#loc3951)
      %35176 = stablehlo.divide %35175, %34047 : tensor<1xf32> loc(#loc3476)
      %35177 = stablehlo.broadcast_in_dim %35176, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3952)
      %35178 = stablehlo.multiply %35177, %34042 : tensor<1x1x2048x2048xf32> loc(#loc3479)
      %35179 = stablehlo.broadcast_in_dim %34036, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %35180 = stablehlo.multiply %35178, %35179 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %35181 = stablehlo.negate %35180 : tensor<1x1x2048x2048xf32> loc(#loc3953)
      %35182 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3954)
      %35183 = stablehlo.reduce(%35181 init: %35182) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3954)
        stablehlo.return %39096 : tensor<f32> loc(#loc3954)
      } loc(#loc3954)
      %35184 = stablehlo.reshape %35183 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3955)
      %35185 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3956)
      %35186 = stablehlo.reduce(%35184 init: %35185) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3956)
        stablehlo.return %39096 : tensor<f32> loc(#loc3956)
      } loc(#loc3956)
      %35187 = stablehlo.broadcast_in_dim %35186, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3957)
      %35188 = func.call @_where_814(%30495, %30097) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35189 = stablehlo.add %35187, %35188 : tensor<1xf32> loc(#loc3797)
      %35190 = stablehlo.divide %35189, %34031 : tensor<1xf32> loc(#loc3476)
      %35191 = stablehlo.broadcast_in_dim %35190, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3952)
      %35192 = stablehlo.broadcast_in_dim %34016, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %35193 = stablehlo.multiply %35191, %35192 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %35194 = stablehlo.add %35180, %35193 : tensor<1x1x2048x2048xf32> loc(#loc3958)
      %35195 = func.call @_where_815(%30495, %30088) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35196 = stablehlo.divide %35195, %34067 : tensor<1xf32> loc(#loc3959)
      %35197 = stablehlo.broadcast_in_dim %35196, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3960)
      %35198 = stablehlo.multiply %35197, %34065 : tensor<1x1x2048x2048xf32> loc(#loc3961)
      %35199 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35200 = stablehlo.broadcast_in_dim %35199, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3962)
      %35201 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35202 = stablehlo.broadcast_in_dim %35201, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3486)
      %35203 = stablehlo.compare  EQ, %34059, %35202,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3486)
      %35204 = stablehlo.select %35203, %35198, %35200 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3963)
      %35205 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35206 = stablehlo.broadcast_in_dim %35205, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3486)
      %35207 = stablehlo.compare  EQ, %34059, %35206,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3486)
      %35208 = stablehlo.select %35207, %35198, %35200 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3963)
      %35209 = stablehlo.negate %35204 : tensor<1x1x2048x2048xf32> loc(#loc3964)
      %35210 = stablehlo.add %35208, %35209 : tensor<1x1x2048x2048xf32> loc(#loc3965)
      %35211 = stablehlo.broadcast_in_dim %34053, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3471)
      %35212 = stablehlo.multiply %35210, %35211 : tensor<1x1x2048x2048xf32> loc(#loc3471)
      %35213 = stablehlo.add %35194, %35212 : tensor<1x1x2048x2048xf32> loc(#loc3958)
      %35214 = stablehlo.convert %35213 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3966)
      %35215 = func.call @_where_816(%30495, %30079) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35216 = stablehlo.multiply %35215, %33797 : tensor<1xf32> loc(#loc3967)
      %35217 = stablehlo.divide %35216, %33792 : tensor<1xf32> loc(#loc3366)
      %35218 = stablehlo.broadcast_in_dim %35217, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3968)
      %35219 = stablehlo.multiply %35218, %33787 : tensor<1x1x2048x2048xf32> loc(#loc3369)
      %35220 = stablehlo.broadcast_in_dim %33781, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %35221 = stablehlo.multiply %35219, %35220 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %35222 = stablehlo.negate %35221 : tensor<1x1x2048x2048xf32> loc(#loc3969)
      %35223 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3970)
      %35224 = stablehlo.reduce(%35222 init: %35223) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3970)
        stablehlo.return %39096 : tensor<f32> loc(#loc3970)
      } loc(#loc3970)
      %35225 = stablehlo.reshape %35224 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3971)
      %35226 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3972)
      %35227 = stablehlo.reduce(%35225 init: %35226) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3972)
        stablehlo.return %39096 : tensor<f32> loc(#loc3972)
      } loc(#loc3972)
      %35228 = stablehlo.broadcast_in_dim %35227, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3973)
      %35229 = func.call @_where_817(%30495, %30070) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35230 = stablehlo.add %35228, %35229 : tensor<1xf32> loc(#loc3797)
      %35231 = stablehlo.divide %35230, %33776 : tensor<1xf32> loc(#loc3366)
      %35232 = stablehlo.broadcast_in_dim %35231, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3968)
      %35233 = stablehlo.broadcast_in_dim %33761, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %35234 = stablehlo.multiply %35232, %35233 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %35235 = stablehlo.add %35221, %35234 : tensor<1x1x2048x2048xf32> loc(#loc3974)
      %35236 = func.call @_where_818(%30495, %30061) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35237 = stablehlo.divide %35236, %33812 : tensor<1xf32> loc(#loc3975)
      %35238 = stablehlo.broadcast_in_dim %35237, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3976)
      %35239 = stablehlo.multiply %35238, %33810 : tensor<1x1x2048x2048xf32> loc(#loc3977)
      %35240 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35241 = stablehlo.broadcast_in_dim %35240, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3978)
      %35242 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35243 = stablehlo.broadcast_in_dim %35242, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3376)
      %35244 = stablehlo.compare  EQ, %33804, %35243,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3376)
      %35245 = stablehlo.select %35244, %35239, %35241 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3979)
      %35246 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35247 = stablehlo.broadcast_in_dim %35246, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3376)
      %35248 = stablehlo.compare  EQ, %33804, %35247,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3376)
      %35249 = stablehlo.select %35248, %35239, %35241 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3979)
      %35250 = stablehlo.negate %35245 : tensor<1x1x2048x2048xf32> loc(#loc3980)
      %35251 = stablehlo.add %35249, %35250 : tensor<1x1x2048x2048xf32> loc(#loc3981)
      %35252 = stablehlo.broadcast_in_dim %33798, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3361)
      %35253 = stablehlo.multiply %35251, %35252 : tensor<1x1x2048x2048xf32> loc(#loc3361)
      %35254 = stablehlo.add %35235, %35253 : tensor<1x1x2048x2048xf32> loc(#loc3974)
      %35255 = stablehlo.convert %35254 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3982)
      %35256 = func.call @_where_819(%30495, %30052) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35257 = stablehlo.multiply %35256, %33542 : tensor<1xf32> loc(#loc3983)
      %35258 = stablehlo.divide %35257, %33537 : tensor<1xf32> loc(#loc3256)
      %35259 = stablehlo.broadcast_in_dim %35258, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3984)
      %35260 = stablehlo.multiply %35259, %33532 : tensor<1x1x2048x2048xf32> loc(#loc3259)
      %35261 = stablehlo.broadcast_in_dim %33526, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %35262 = stablehlo.multiply %35260, %35261 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %35263 = stablehlo.negate %35262 : tensor<1x1x2048x2048xf32> loc(#loc3985)
      %35264 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3986)
      %35265 = stablehlo.reduce(%35263 init: %35264) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3986)
        stablehlo.return %39096 : tensor<f32> loc(#loc3986)
      } loc(#loc3986)
      %35266 = stablehlo.reshape %35265 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc3987)
      %35267 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3988)
      %35268 = stablehlo.reduce(%35266 init: %35267) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3988)
        stablehlo.return %39096 : tensor<f32> loc(#loc3988)
      } loc(#loc3988)
      %35269 = stablehlo.broadcast_in_dim %35268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc3989)
      %35270 = func.call @_where_820(%30495, %30043) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35271 = stablehlo.add %35269, %35270 : tensor<1xf32> loc(#loc3797)
      %35272 = stablehlo.divide %35271, %33521 : tensor<1xf32> loc(#loc3256)
      %35273 = stablehlo.broadcast_in_dim %35272, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3984)
      %35274 = stablehlo.broadcast_in_dim %33506, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %35275 = stablehlo.multiply %35273, %35274 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %35276 = stablehlo.add %35262, %35275 : tensor<1x1x2048x2048xf32> loc(#loc3990)
      %35277 = func.call @_where_821(%30495, %30034) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35278 = stablehlo.divide %35277, %33557 : tensor<1xf32> loc(#loc3991)
      %35279 = stablehlo.broadcast_in_dim %35278, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3992)
      %35280 = stablehlo.multiply %35279, %33555 : tensor<1x1x2048x2048xf32> loc(#loc3993)
      %35281 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35282 = stablehlo.broadcast_in_dim %35281, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc3994)
      %35283 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35284 = stablehlo.broadcast_in_dim %35283, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3266)
      %35285 = stablehlo.compare  EQ, %33549, %35284,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3266)
      %35286 = stablehlo.select %35285, %35280, %35282 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3995)
      %35287 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35288 = stablehlo.broadcast_in_dim %35287, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3266)
      %35289 = stablehlo.compare  EQ, %33549, %35288,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3266)
      %35290 = stablehlo.select %35289, %35280, %35282 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc3995)
      %35291 = stablehlo.negate %35286 : tensor<1x1x2048x2048xf32> loc(#loc3996)
      %35292 = stablehlo.add %35290, %35291 : tensor<1x1x2048x2048xf32> loc(#loc3997)
      %35293 = stablehlo.broadcast_in_dim %33543, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3251)
      %35294 = stablehlo.multiply %35292, %35293 : tensor<1x1x2048x2048xf32> loc(#loc3251)
      %35295 = stablehlo.add %35276, %35294 : tensor<1x1x2048x2048xf32> loc(#loc3990)
      %35296 = stablehlo.convert %35295 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3998)
      %35297 = func.call @_where_822(%30495, %30025) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35298 = stablehlo.multiply %35297, %33287 : tensor<1xf32> loc(#loc3999)
      %35299 = stablehlo.divide %35298, %33282 : tensor<1xf32> loc(#loc3146)
      %35300 = stablehlo.broadcast_in_dim %35299, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4000)
      %35301 = stablehlo.multiply %35300, %33277 : tensor<1x1x2048x2048xf32> loc(#loc3149)
      %35302 = stablehlo.broadcast_in_dim %33271, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %35303 = stablehlo.multiply %35301, %35302 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %35304 = stablehlo.negate %35303 : tensor<1x1x2048x2048xf32> loc(#loc4001)
      %35305 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4002)
      %35306 = stablehlo.reduce(%35304 init: %35305) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4002)
        stablehlo.return %39096 : tensor<f32> loc(#loc4002)
      } loc(#loc4002)
      %35307 = stablehlo.reshape %35306 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc4003)
      %35308 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4004)
      %35309 = stablehlo.reduce(%35307 init: %35308) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4004)
        stablehlo.return %39096 : tensor<f32> loc(#loc4004)
      } loc(#loc4004)
      %35310 = stablehlo.broadcast_in_dim %35309, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc4005)
      %35311 = func.call @_where_823(%30495, %30016) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35312 = stablehlo.add %35310, %35311 : tensor<1xf32> loc(#loc3797)
      %35313 = stablehlo.divide %35312, %33266 : tensor<1xf32> loc(#loc3146)
      %35314 = stablehlo.broadcast_in_dim %35313, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4000)
      %35315 = stablehlo.broadcast_in_dim %33251, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %35316 = stablehlo.multiply %35314, %35315 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %35317 = stablehlo.add %35303, %35316 : tensor<1x1x2048x2048xf32> loc(#loc4006)
      %35318 = func.call @_where_824(%30495, %30007) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35319 = stablehlo.divide %35318, %33302 : tensor<1xf32> loc(#loc4007)
      %35320 = stablehlo.broadcast_in_dim %35319, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4008)
      %35321 = stablehlo.multiply %35320, %33300 : tensor<1x1x2048x2048xf32> loc(#loc4009)
      %35322 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35323 = stablehlo.broadcast_in_dim %35322, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc4010)
      %35324 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35325 = stablehlo.broadcast_in_dim %35324, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3156)
      %35326 = stablehlo.compare  EQ, %33294, %35325,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3156)
      %35327 = stablehlo.select %35326, %35321, %35323 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc4011)
      %35328 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35329 = stablehlo.broadcast_in_dim %35328, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc3156)
      %35330 = stablehlo.compare  EQ, %33294, %35329,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc3156)
      %35331 = stablehlo.select %35330, %35321, %35323 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc4011)
      %35332 = stablehlo.negate %35327 : tensor<1x1x2048x2048xf32> loc(#loc4012)
      %35333 = stablehlo.add %35331, %35332 : tensor<1x1x2048x2048xf32> loc(#loc4013)
      %35334 = stablehlo.broadcast_in_dim %33288, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc3141)
      %35335 = stablehlo.multiply %35333, %35334 : tensor<1x1x2048x2048xf32> loc(#loc3141)
      %35336 = stablehlo.add %35317, %35335 : tensor<1x1x2048x2048xf32> loc(#loc4006)
      %35337 = stablehlo.convert %35336 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc4014)
      %35338 = func.call @_where_825(%30495, %29998) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35339 = stablehlo.multiply %35338, %30992 : tensor<1xf32> loc(#loc4015)
      %35340 = stablehlo.divide %35339, %30987 : tensor<1xf32> loc(#loc2156)
      %35341 = stablehlo.broadcast_in_dim %35340, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4016)
      %35342 = stablehlo.multiply %35341, %30982 : tensor<1x1x2048x2048xf32> loc(#loc2159)
      %35343 = stablehlo.broadcast_in_dim %30976, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %35344 = stablehlo.multiply %35342, %35343 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %35345 = stablehlo.negate %35344 : tensor<1x1x2048x2048xf32> loc(#loc4017)
      %35346 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4018)
      %35347 = stablehlo.reduce(%35345 init: %35346) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4018)
        stablehlo.return %39096 : tensor<f32> loc(#loc4018)
      } loc(#loc4018)
      %35348 = stablehlo.reshape %35347 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc4019)
      %35349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4020)
      %35350 = stablehlo.reduce(%35348 init: %35349) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4020)
        stablehlo.return %39096 : tensor<f32> loc(#loc4020)
      } loc(#loc4020)
      %35351 = stablehlo.broadcast_in_dim %35350, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc4021)
      %35352 = func.call @_where_826(%30495, %29989) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35353 = stablehlo.add %35351, %35352 : tensor<1xf32> loc(#loc3797)
      %35354 = stablehlo.divide %35353, %30971 : tensor<1xf32> loc(#loc2156)
      %35355 = stablehlo.broadcast_in_dim %35354, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4016)
      %35356 = stablehlo.broadcast_in_dim %30956, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %35357 = stablehlo.multiply %35355, %35356 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %35358 = stablehlo.add %35344, %35357 : tensor<1x1x2048x2048xf32> loc(#loc4022)
      %35359 = func.call @_where_827(%30495, %29980) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35360 = stablehlo.divide %35359, %31007 : tensor<1xf32> loc(#loc4023)
      %35361 = stablehlo.broadcast_in_dim %35360, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4024)
      %35362 = stablehlo.multiply %35361, %31005 : tensor<1x1x2048x2048xf32> loc(#loc4025)
      %35363 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35364 = stablehlo.broadcast_in_dim %35363, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc4026)
      %35365 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35366 = stablehlo.broadcast_in_dim %35365, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2166)
      %35367 = stablehlo.compare  EQ, %30999, %35366,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2166)
      %35368 = stablehlo.select %35367, %35362, %35364 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc4027)
      %35369 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35370 = stablehlo.broadcast_in_dim %35369, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2166)
      %35371 = stablehlo.compare  EQ, %30999, %35370,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2166)
      %35372 = stablehlo.select %35371, %35362, %35364 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc4027)
      %35373 = stablehlo.negate %35368 : tensor<1x1x2048x2048xf32> loc(#loc4028)
      %35374 = stablehlo.add %35372, %35373 : tensor<1x1x2048x2048xf32> loc(#loc4029)
      %35375 = stablehlo.broadcast_in_dim %30993, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2151)
      %35376 = stablehlo.multiply %35374, %35375 : tensor<1x1x2048x2048xf32> loc(#loc2151)
      %35377 = stablehlo.add %35358, %35376 : tensor<1x1x2048x2048xf32> loc(#loc4022)
      %35378 = stablehlo.convert %35377 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc4030)
      %35379 = func.call @_where_828(%30495, %29971) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35380 = stablehlo.multiply %35379, %30737 : tensor<1xf32> loc(#loc4031)
      %35381 = stablehlo.divide %35380, %30732 : tensor<1xf32> loc(#loc2046)
      %35382 = stablehlo.broadcast_in_dim %35381, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4032)
      %35383 = stablehlo.multiply %35382, %30727 : tensor<1x1x2048x2048xf32> loc(#loc2049)
      %35384 = stablehlo.broadcast_in_dim %30721, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %35385 = stablehlo.multiply %35383, %35384 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %35386 = stablehlo.negate %35385 : tensor<1x1x2048x2048xf32> loc(#loc4033)
      %35387 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4034)
      %35388 = stablehlo.reduce(%35386 init: %35387) across dimensions = [2, 3] : (tensor<1x1x2048x2048xf32>, tensor<f32>) -> tensor<1x1xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4034)
        stablehlo.return %39096 : tensor<f32> loc(#loc4034)
      } loc(#loc4034)
      %35389 = stablehlo.reshape %35388 : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32> loc(#loc4035)
      %35390 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4036)
      %35391 = stablehlo.reduce(%35389 init: %35390) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>, tensor<f32>) -> tensor<f32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4036)
        stablehlo.return %39096 : tensor<f32> loc(#loc4036)
      } loc(#loc4036)
      %35392 = stablehlo.broadcast_in_dim %35391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc4037)
      %35393 = func.call @_where_829(%30495, %29962) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35394 = stablehlo.add %35392, %35393 : tensor<1xf32> loc(#loc3797)
      %35395 = stablehlo.divide %35394, %30716 : tensor<1xf32> loc(#loc2046)
      %35396 = stablehlo.broadcast_in_dim %35395, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4032)
      %35397 = stablehlo.broadcast_in_dim %30701, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %35398 = stablehlo.multiply %35396, %35397 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %35399 = stablehlo.add %35385, %35398 : tensor<1x1x2048x2048xf32> loc(#loc4038)
      %35400 = func.call @_where_830(%30495, %29953) : (tensor<1xi1>, tensor<1xf32>) -> tensor<1xf32> loc(#loc3789)
      %35401 = stablehlo.divide %35400, %30752 : tensor<1xf32> loc(#loc4039)
      %35402 = stablehlo.broadcast_in_dim %35401, dims = [0] : (tensor<1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4040)
      %35403 = stablehlo.multiply %35402, %30750 : tensor<1x1x2048x2048xf32> loc(#loc4041)
      %35404 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc1984)
      %35405 = stablehlo.broadcast_in_dim %35404, dims = [] : (tensor<f32>) -> tensor<1x1x2048x2048xf32> loc(#loc4042)
      %35406 = stablehlo.constant dense<false> : tensor<i1> loc(#loc1984)
      %35407 = stablehlo.broadcast_in_dim %35406, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2056)
      %35408 = stablehlo.compare  EQ, %30744, %35407,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2056)
      %35409 = stablehlo.select %35408, %35403, %35405 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc4043)
      %35410 = stablehlo.constant dense<true> : tensor<i1> loc(#loc1984)
      %35411 = stablehlo.broadcast_in_dim %35410, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc2056)
      %35412 = stablehlo.compare  EQ, %30744, %35411,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc2056)
      %35413 = stablehlo.select %35412, %35403, %35405 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xf32> loc(#loc4043)
      %35414 = stablehlo.negate %35409 : tensor<1x1x2048x2048xf32> loc(#loc4044)
      %35415 = stablehlo.add %35413, %35414 : tensor<1x1x2048x2048xf32> loc(#loc4045)
      %35416 = stablehlo.broadcast_in_dim %30738, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc2041)
      %35417 = stablehlo.multiply %35415, %35416 : tensor<1x1x2048x2048xf32> loc(#loc2041)
      %35418 = stablehlo.add %35399, %35417 : tensor<1x1x2048x2048xf32> loc(#loc4038)
      %35419 = stablehlo.convert %35418 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc4046)
      %35420 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35421 = stablehlo.pad %iterArg_406, %35420, low = [0, 0, 0, 0], high = [1, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc4047)
      %35422 = func.call @_pad_831(%35421) : (tensor<2x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc4048)
      %35423 = stablehlo.add %29944, %35422 : tensor<1x1x2048x2048xbf16> loc(#loc4049)
      %35424 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %35425 = stablehlo.broadcast_in_dim %35424, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc4050)
      %35426 = stablehlo.multiply %35423, %35425 : tensor<1x1x2048x2048xbf16> loc(#loc4050)
      %35427 = stablehlo.broadcast_in_dim %34763, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc4051)
      %35428 = stablehlo.multiply %35426, %35427 : tensor<1x1x2048x2048xbf16> loc(#loc4051)
      %35429 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4052)
      %35430 = stablehlo.reduce(%35428 init: %35429) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4052)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4052)
      } loc(#loc4052)
      %35431 = stablehlo.reshape %35430 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4053)
      %35432 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4054)
      %35433 = stablehlo.reduce(%35431 init: %35432) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4054)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4054)
      } loc(#loc4054)
      %35434 = stablehlo.broadcast_in_dim %35433, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4055)
      %35435 = func.call @_where_832(%30578, %35434) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35436 = arith.select %ptrue, %35428, %35428 : tensor<1x1x2048x2048xbf16>
      %35437:2 = func.call @_einsum_833(%30579#0, %34758, %35436) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4058)
      %35438 = func.call @_where_834(%30579#1, %35437#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %35439 = stablehlo.broadcast_in_dim %34756, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3787)
      %35440 = stablehlo.multiply %35437#0, %35439 : tensor<1x1x2048x8192xbf16> loc(#loc3787)
      %35441 = stablehlo.multiply %34727, %35440 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35442 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %35443 = stablehlo.broadcast_in_dim %35442, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35444 = stablehlo.multiply %35443, %35441 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35445 = stablehlo.multiply %35444, %34744 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35446 = stablehlo.multiply %35445, %34741 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35447 = stablehlo.add %35445, %35446 : tensor<1x1x2048x8192xbf16> loc(#loc4059)
      %35448 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %35449 = stablehlo.broadcast_in_dim %35448, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35450 = stablehlo.multiply %35449, %35447 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35451 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %35452 = stablehlo.broadcast_in_dim %35451, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35453 = stablehlo.multiply %35452, %35450 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35454 = stablehlo.multiply %35453, %34733 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35455 = stablehlo.add %35450, %35454 : tensor<1x1x2048x8192xbf16> loc(#loc4059)
      %35456 = stablehlo.multiply %35440, %34750 : tensor<1x1x2048x8192xbf16> loc(#loc3780)
      %35457 = stablehlo.add %35455, %35456 : tensor<1x1x2048x8192xbf16> loc(#loc4059)
      %35458 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4060)
      %35459 = stablehlo.reduce(%35457 init: %35458) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4060)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4060)
      } loc(#loc4060)
      %35460 = stablehlo.reshape %35459 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4061)
      %35461 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4062)
      %35462 = stablehlo.reduce(%35460 init: %35461) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4062)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4062)
      } loc(#loc4062)
      %35463 = stablehlo.broadcast_in_dim %35462, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4063)
      %35464 = func.call @_where_835(%30576#1, %35463) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %35465 = arith.select %ptrue, %35457, %35457 : tensor<1x1x2048x8192xbf16>
      %35466:2 = func.call @_einsum_836(%30577#0, %34722, %35465) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4064)
      %35467 = func.call @_where_837(%30577#1, %35466#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %35468 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4065)
      %35469 = stablehlo.reduce(%35466#0 init: %35468) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4065)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4065)
      } loc(#loc4065)
      %35470 = stablehlo.reshape %35469 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4066)
      %35471 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4067)
      %35472 = stablehlo.reduce(%35470 init: %35471) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4067)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4067)
      } loc(#loc4067)
      %35473 = stablehlo.broadcast_in_dim %35472, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4068)
      %35474 = func.call @_where_838(%30580#1, %35473) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35475 = stablehlo.multiply %34713, %35466#0 : tensor<1x1x2048x2048xbf16> loc(#loc3773)
      %35476 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4069)
      %35477 = stablehlo.reduce(%35475 init: %35476) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4069)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4069)
      } loc(#loc4069)
      %35478 = stablehlo.reshape %35477 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4070)
      %35479 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4071)
      %35480 = stablehlo.reduce(%35478 init: %35479) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4071)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4071)
      } loc(#loc4071)
      %35481 = stablehlo.broadcast_in_dim %35480, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4072)
      %35482 = func.call @_where_839(%30581#1, %35481) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35483 = stablehlo.broadcast_in_dim %34717, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3773)
      %35484 = stablehlo.multiply %35466#0, %35483 : tensor<1x1x2048x2048xbf16> loc(#loc3773)
      %35485 = stablehlo.multiply %34703, %35484 : tensor<1x1x2048x2048xbf16> loc(#loc3770)
      %35486 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4073)
      %35487 = stablehlo.reduce(%35485 init: %35486) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4073)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4073)
      } loc(#loc4073)
      %35488 = stablehlo.reshape %35487 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4074)
      %35489 = stablehlo.multiply %35488, %34711 : tensor<1x1x2048x1xbf16> loc(#loc3770)
      %35490 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35491 = stablehlo.broadcast_in_dim %35490, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3764)
      %35492 = stablehlo.divide %35489, %35491 : tensor<1x1x2048x1xbf16> loc(#loc3764)
      %35493 = stablehlo.convert %35492 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3760)
      %35494 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4075)
      %35495 = stablehlo.reduce(%35493 init: %35494) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4075)
        stablehlo.return %39096 : tensor<f32> loc(#loc4075)
      } loc(#loc4075)
      %35496 = stablehlo.broadcast_in_dim %35495, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4076)
      %35497 = stablehlo.broadcast_in_dim %35496, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4077)
      %35498 = stablehlo.convert %35497 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3763)
      %35499 = stablehlo.multiply %35498, %34693 : tensor<1x1x2048x2048xbf16> loc(#loc3766)
      %35500 = stablehlo.negate %35499 : tensor<1x1x2048x2048xbf16> loc(#loc4078)
      %35501 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4079)
      %35502 = stablehlo.reduce(%35500 init: %35501) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4079)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4079)
      } loc(#loc4079)
      %35503 = stablehlo.reshape %35502 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4080)
      %35504 = stablehlo.add %35423, %35499 : tensor<1x1x2048x2048xbf16> loc(#loc4081)
      %35505 = stablehlo.broadcast_in_dim %34707, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3770)
      %35506 = stablehlo.multiply %35484, %35505 : tensor<1x1x2048x2048xbf16> loc(#loc3770)
      %35507 = stablehlo.negate %35506 : tensor<1x1x2048x2048xbf16> loc(#loc4078)
      %35508 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4079)
      %35509 = stablehlo.reduce(%35507 init: %35508) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4079)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4079)
      } loc(#loc4079)
      %35510 = stablehlo.reshape %35509 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4080)
      %35511 = stablehlo.add %35504, %35506 : tensor<1x1x2048x2048xbf16> loc(#loc4081)
      %35512 = stablehlo.add %35503, %35510 : tensor<1x1x2048x1xbf16> loc(#loc4081)
      %35513 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35514 = stablehlo.broadcast_in_dim %35513, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3764)
      %35515 = stablehlo.divide %35512, %35514 : tensor<1x1x2048x1xbf16> loc(#loc3764)
      %35516 = stablehlo.convert %35515 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3760)
      %35517 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4075)
      %35518 = stablehlo.reduce(%35516 init: %35517) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4075)
        stablehlo.return %39096 : tensor<f32> loc(#loc4075)
      } loc(#loc4075)
      %35519 = stablehlo.broadcast_in_dim %35518, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4076)
      %35520 = stablehlo.broadcast_in_dim %35519, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4077)
      %35521 = stablehlo.convert %35520 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3763)
      %35522 = stablehlo.add %35511, %35521 : tensor<1x1x2048x2048xbf16> loc(#loc4082)
      %35523 = stablehlo.add %35132, %35522 : tensor<1x1x2048x2048xbf16> loc(#loc4083)
      %35524 = arith.select %ptrue, %35522, %35522 : tensor<1x1x2048x2048xbf16>
      %35525:2 = func.call @_einsum_840(%30586#0, %34675, %35524) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4084)
      %35526 = func.call @_where_841(%30586#1, %35525#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %35527 = arith.select %ptrue, %35525#0, %35525#0 : tensor<1x1x2048x32x64xbf16>
      %35528:2 = func.call @_einsum_842(%34673, %34630, %35527) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4085)
      %35529 = arith.select %ptrue, %35528#1, %35528#1 : tensor<1x1x2048x32x64xbf16>
      %35530 = stablehlo.broadcast_in_dim %35529, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4086)
      %35531 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35532 = stablehlo.pad %35530, %35531, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4087)
      %35533 = stablehlo.convert %35528#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4088)
      %35534 = stablehlo.broadcast_in_dim %34672, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3751)
      %35535 = stablehlo.multiply %35533, %35534 : tensor<1x1x32x2048x2048xf32> loc(#loc3751)
      %35536 = stablehlo.multiply %35535, %34663 : tensor<1x1x32x2048x2048xf32> loc(#loc3751)
      %35537 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3749)
      %35538 = stablehlo.reduce(%35536 init: %35537) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3749)
        stablehlo.return %39096 : tensor<f32> loc(#loc3749)
      } loc(#loc3749)
      %35539 = stablehlo.reshape %35538 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4089)
      %35540 = stablehlo.negate %35539 : tensor<1x1x32x2048x1xf32> loc(#loc4090)
      %35541 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4091)
      %35542 = stablehlo.reduce(%35540 init: %35541) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4091)
        stablehlo.return %39096 : tensor<f32> loc(#loc4091)
      } loc(#loc4091)
      %35543 = stablehlo.broadcast_in_dim %35542, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4092)
      %35544 = stablehlo.broadcast_in_dim %35543, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4093)
      %35545 = stablehlo.broadcast_in_dim %34666, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3750)
      %35546 = stablehlo.divide %35533, %35545 : tensor<1x1x32x2048x2048xf32> loc(#loc3750)
      %35547 = stablehlo.add %35544, %35546 : tensor<1x1x32x2048x2048xf32> loc(#loc4094)
      %35548 = stablehlo.multiply %35547, %34663 : tensor<1x1x32x2048x2048xf32> loc(#loc3751)
      %35549 = func.call @_where_843(%34657#1, %35548) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4095)
      %35550 = stablehlo.convert %35549 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4096)
      %35551 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %35552 = stablehlo.broadcast_in_dim %35551, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3739)
      %35553 = stablehlo.multiply %35552, %35550 : tensor<1x1x32x2048x2048xbf16> loc(#loc3739)
      %35554 = stablehlo.multiply %35553, %34645 : tensor<1x1x32x2048x2048xbf16> loc(#loc3739)
      %35555 = stablehlo.multiply %35554, %34642 : tensor<1x1x32x2048x2048xbf16> loc(#loc3739)
      %35556 = stablehlo.add %35554, %35555 : tensor<1x1x32x2048x2048xbf16> loc(#loc4097)
      %35557 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %35558 = stablehlo.broadcast_in_dim %35557, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3736)
      %35559 = stablehlo.divide %35556, %35558 : tensor<1x1x32x2048x2048xbf16> loc(#loc3736)
      %35560:2 = func.call @_einsum_844(%34629, %34637, %35559) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4098)
      %35561 = arith.select %ptrue, %35560#1, %35560#1 : tensor<1x1x2048x32x64xbf16>
      %35562 = stablehlo.broadcast_in_dim %35561, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4086)
      %35563 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35564 = stablehlo.pad %35562, %35563, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4099)
      %35565 = stablehlo.add %35532, %35564 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4100)
      %35566 = stablehlo.multiply %34628, %35560#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3734)
      %35567 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4101)
      %35568 = stablehlo.reduce(%35566 init: %35567) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4101)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4101)
      } loc(#loc4101)
      %35569 = stablehlo.reshape %35568 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4102)
      %35570 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4103)
      %35571 = stablehlo.reduce(%35569 init: %35570) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4103)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4103)
      } loc(#loc4103)
      %35572 = stablehlo.broadcast_in_dim %35571, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4104)
      %35573 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %35574 = stablehlo.broadcast_in_dim %35573, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3732)
      %35575 = stablehlo.multiply %35574, %35572 : tensor<1x64xbf16> loc(#loc3732)
      %35576 = func.call @softplus_845(%34631#1, %35575) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4105)
      %35577 = func.call @_where_846(%30585#1, %35576) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %35578 = stablehlo.broadcast_in_dim %34635, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3734)
      %35579 = stablehlo.multiply %35560#0, %35578 : tensor<1x1x2048x32x64xbf16> loc(#loc3734)
      %35580 = arith.select %ptrue, %35579, %35579 : tensor<1x1x2048x32x64xbf16>
      %35581 = stablehlo.broadcast_in_dim %35580, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4086)
      %35582 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35583 = stablehlo.pad %35581, %35582, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4106)
      %35584 = stablehlo.add %35565, %35583 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4100)
      %35585:2 = func.call @_einsum_847(%34620, %30584#0, %35584) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4107)
      %35586 = func.call @_where_848(%30584#1, %35585#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %35587 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4108)
      %35588 = stablehlo.reduce(%35585#0 init: %35587) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4108)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4108)
      } loc(#loc4108)
      %35589 = stablehlo.reshape %35588 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4109)
      %35590 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4110)
      %35591 = stablehlo.reduce(%35589 init: %35590) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4110)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4110)
      } loc(#loc4110)
      %35592 = stablehlo.broadcast_in_dim %35591, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4111)
      %35593 = func.call @_where_849(%30582#1, %35592) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35594 = stablehlo.multiply %34611, %35585#0 : tensor<1x1x2048x2048xbf16> loc(#loc3722)
      %35595 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4112)
      %35596 = stablehlo.reduce(%35594 init: %35595) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4112)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4112)
      } loc(#loc4112)
      %35597 = stablehlo.reshape %35596 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4113)
      %35598 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4114)
      %35599 = stablehlo.reduce(%35597 init: %35598) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4114)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4114)
      } loc(#loc4114)
      %35600 = stablehlo.broadcast_in_dim %35599, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4115)
      %35601 = func.call @_where_850(%30583#1, %35600) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35602 = stablehlo.broadcast_in_dim %34615, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3722)
      %35603 = stablehlo.multiply %35585#0, %35602 : tensor<1x1x2048x2048xbf16> loc(#loc3722)
      %35604 = stablehlo.multiply %34601, %35603 : tensor<1x1x2048x2048xbf16> loc(#loc3719)
      %35605 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4116)
      %35606 = stablehlo.reduce(%35604 init: %35605) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4116)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4116)
      } loc(#loc4116)
      %35607 = stablehlo.reshape %35606 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4117)
      %35608 = stablehlo.multiply %35607, %34609 : tensor<1x1x2048x1xbf16> loc(#loc3719)
      %35609 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35610 = stablehlo.broadcast_in_dim %35609, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3713)
      %35611 = stablehlo.divide %35608, %35610 : tensor<1x1x2048x1xbf16> loc(#loc3713)
      %35612 = stablehlo.convert %35611 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3709)
      %35613 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4118)
      %35614 = stablehlo.reduce(%35612 init: %35613) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4118)
        stablehlo.return %39096 : tensor<f32> loc(#loc4118)
      } loc(#loc4118)
      %35615 = stablehlo.broadcast_in_dim %35614, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4119)
      %35616 = stablehlo.broadcast_in_dim %35615, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4120)
      %35617 = stablehlo.convert %35616 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3712)
      %35618 = stablehlo.multiply %35617, %34591 : tensor<1x1x2048x2048xbf16> loc(#loc3715)
      %35619 = stablehlo.negate %35618 : tensor<1x1x2048x2048xbf16> loc(#loc4121)
      %35620 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4122)
      %35621 = stablehlo.reduce(%35619 init: %35620) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4122)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4122)
      } loc(#loc4122)
      %35622 = stablehlo.reshape %35621 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4123)
      %35623 = stablehlo.add %35523, %35618 : tensor<1x1x2048x2048xbf16> loc(#loc4124)
      %35624 = stablehlo.broadcast_in_dim %34605, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3719)
      %35625 = stablehlo.multiply %35603, %35624 : tensor<1x1x2048x2048xbf16> loc(#loc3719)
      %35626 = stablehlo.negate %35625 : tensor<1x1x2048x2048xbf16> loc(#loc4121)
      %35627 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4122)
      %35628 = stablehlo.reduce(%35626 init: %35627) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4122)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4122)
      } loc(#loc4122)
      %35629 = stablehlo.reshape %35628 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4123)
      %35630 = stablehlo.add %35623, %35625 : tensor<1x1x2048x2048xbf16> loc(#loc4124)
      %35631 = stablehlo.add %35622, %35629 : tensor<1x1x2048x1xbf16> loc(#loc4124)
      %35632 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35633 = stablehlo.broadcast_in_dim %35632, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3713)
      %35634 = stablehlo.divide %35631, %35633 : tensor<1x1x2048x1xbf16> loc(#loc3713)
      %35635 = stablehlo.convert %35634 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3709)
      %35636 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4118)
      %35637 = stablehlo.reduce(%35635 init: %35636) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4118)
        stablehlo.return %39096 : tensor<f32> loc(#loc4118)
      } loc(#loc4118)
      %35638 = stablehlo.broadcast_in_dim %35637, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4119)
      %35639 = stablehlo.broadcast_in_dim %35638, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4120)
      %35640 = stablehlo.convert %35639 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3712)
      %35641 = stablehlo.add %35630, %35640 : tensor<1x1x2048x2048xbf16> loc(#loc4125)
      %35642 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %35643 = stablehlo.broadcast_in_dim %35642, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3684)
      %35644 = stablehlo.multiply %35641, %35643 : tensor<1x1x2048x2048xbf16> loc(#loc3684)
      %35645 = stablehlo.broadcast_in_dim %34513, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3683)
      %35646 = stablehlo.multiply %35644, %35645 : tensor<1x1x2048x2048xbf16> loc(#loc3683)
      %35647 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4126)
      %35648 = stablehlo.reduce(%35646 init: %35647) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4126)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4126)
      } loc(#loc4126)
      %35649 = stablehlo.reshape %35648 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4127)
      %35650 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4128)
      %35651 = stablehlo.reduce(%35649 init: %35650) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4128)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4128)
      } loc(#loc4128)
      %35652 = stablehlo.broadcast_in_dim %35651, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4129)
      %35653 = func.call @_where_851(%30567#1, %35652) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35654 = arith.select %ptrue, %35646, %35646 : tensor<1x1x2048x2048xbf16>
      %35655:2 = func.call @_einsum_852(%30568#0, %34503, %35654) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4130)
      %35656 = func.call @_where_853(%30568#1, %35655#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %35657 = stablehlo.broadcast_in_dim %34501, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3677)
      %35658 = stablehlo.multiply %35655#0, %35657 : tensor<1x1x2048x8192xbf16> loc(#loc3677)
      %35659 = stablehlo.multiply %34472, %35658 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35660 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %35661 = stablehlo.broadcast_in_dim %35660, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35662 = stablehlo.multiply %35661, %35659 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35663 = stablehlo.multiply %35662, %34489 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35664 = stablehlo.multiply %35663, %34486 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35665 = stablehlo.add %35663, %35664 : tensor<1x1x2048x8192xbf16> loc(#loc4131)
      %35666 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %35667 = stablehlo.broadcast_in_dim %35666, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35668 = stablehlo.multiply %35667, %35665 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35669 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %35670 = stablehlo.broadcast_in_dim %35669, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35671 = stablehlo.multiply %35670, %35668 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35672 = stablehlo.multiply %35671, %34478 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35673 = stablehlo.add %35668, %35672 : tensor<1x1x2048x8192xbf16> loc(#loc4131)
      %35674 = stablehlo.multiply %35658, %34495 : tensor<1x1x2048x8192xbf16> loc(#loc3670)
      %35675 = stablehlo.add %35673, %35674 : tensor<1x1x2048x8192xbf16> loc(#loc4131)
      %35676 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4132)
      %35677 = stablehlo.reduce(%35675 init: %35676) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4132)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4132)
      } loc(#loc4132)
      %35678 = stablehlo.reshape %35677 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4133)
      %35679 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4134)
      %35680 = stablehlo.reduce(%35678 init: %35679) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4134)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4134)
      } loc(#loc4134)
      %35681 = stablehlo.broadcast_in_dim %35680, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4135)
      %35682 = func.call @_where_854(%30565#1, %35681) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %35683 = arith.select %ptrue, %35675, %35675 : tensor<1x1x2048x8192xbf16>
      %35684:2 = func.call @_einsum_855(%30566#0, %34467, %35683) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4136)
      %35685 = func.call @_where_856(%30566#1, %35684#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %35686 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4137)
      %35687 = stablehlo.reduce(%35684#0 init: %35686) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4137)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4137)
      } loc(#loc4137)
      %35688 = stablehlo.reshape %35687 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4138)
      %35689 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4139)
      %35690 = stablehlo.reduce(%35688 init: %35689) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4139)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4139)
      } loc(#loc4139)
      %35691 = stablehlo.broadcast_in_dim %35690, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4140)
      %35692 = func.call @_where_857(%30569#1, %35691) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35693 = stablehlo.multiply %34458, %35684#0 : tensor<1x1x2048x2048xbf16> loc(#loc3663)
      %35694 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4141)
      %35695 = stablehlo.reduce(%35693 init: %35694) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4141)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4141)
      } loc(#loc4141)
      %35696 = stablehlo.reshape %35695 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4142)
      %35697 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4143)
      %35698 = stablehlo.reduce(%35696 init: %35697) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4143)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4143)
      } loc(#loc4143)
      %35699 = stablehlo.broadcast_in_dim %35698, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4144)
      %35700 = func.call @_where_858(%30570#1, %35699) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35701 = stablehlo.broadcast_in_dim %34462, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3663)
      %35702 = stablehlo.multiply %35684#0, %35701 : tensor<1x1x2048x2048xbf16> loc(#loc3663)
      %35703 = stablehlo.multiply %34448, %35702 : tensor<1x1x2048x2048xbf16> loc(#loc3660)
      %35704 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4145)
      %35705 = stablehlo.reduce(%35703 init: %35704) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4145)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4145)
      } loc(#loc4145)
      %35706 = stablehlo.reshape %35705 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4146)
      %35707 = stablehlo.multiply %35706, %34456 : tensor<1x1x2048x1xbf16> loc(#loc3660)
      %35708 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35709 = stablehlo.broadcast_in_dim %35708, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3654)
      %35710 = stablehlo.divide %35707, %35709 : tensor<1x1x2048x1xbf16> loc(#loc3654)
      %35711 = stablehlo.convert %35710 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3650)
      %35712 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4147)
      %35713 = stablehlo.reduce(%35711 init: %35712) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4147)
        stablehlo.return %39096 : tensor<f32> loc(#loc4147)
      } loc(#loc4147)
      %35714 = stablehlo.broadcast_in_dim %35713, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4148)
      %35715 = stablehlo.broadcast_in_dim %35714, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4149)
      %35716 = stablehlo.convert %35715 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3653)
      %35717 = stablehlo.multiply %35716, %34438 : tensor<1x1x2048x2048xbf16> loc(#loc3656)
      %35718 = stablehlo.negate %35717 : tensor<1x1x2048x2048xbf16> loc(#loc4150)
      %35719 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4151)
      %35720 = stablehlo.reduce(%35718 init: %35719) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4151)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4151)
      } loc(#loc4151)
      %35721 = stablehlo.reshape %35720 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4152)
      %35722 = stablehlo.add %35641, %35717 : tensor<1x1x2048x2048xbf16> loc(#loc4153)
      %35723 = stablehlo.broadcast_in_dim %34452, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3660)
      %35724 = stablehlo.multiply %35702, %35723 : tensor<1x1x2048x2048xbf16> loc(#loc3660)
      %35725 = stablehlo.negate %35724 : tensor<1x1x2048x2048xbf16> loc(#loc4150)
      %35726 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4151)
      %35727 = stablehlo.reduce(%35725 init: %35726) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4151)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4151)
      } loc(#loc4151)
      %35728 = stablehlo.reshape %35727 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4152)
      %35729 = stablehlo.add %35722, %35724 : tensor<1x1x2048x2048xbf16> loc(#loc4153)
      %35730 = stablehlo.add %35721, %35728 : tensor<1x1x2048x1xbf16> loc(#loc4153)
      %35731 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35732 = stablehlo.broadcast_in_dim %35731, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3654)
      %35733 = stablehlo.divide %35730, %35732 : tensor<1x1x2048x1xbf16> loc(#loc3654)
      %35734 = stablehlo.convert %35733 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3650)
      %35735 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4147)
      %35736 = stablehlo.reduce(%35734 init: %35735) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4147)
        stablehlo.return %39096 : tensor<f32> loc(#loc4147)
      } loc(#loc4147)
      %35737 = stablehlo.broadcast_in_dim %35736, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4148)
      %35738 = stablehlo.broadcast_in_dim %35737, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4149)
      %35739 = stablehlo.convert %35738 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3653)
      %35740 = stablehlo.add %35729, %35739 : tensor<1x1x2048x2048xbf16> loc(#loc4154)
      %35741 = stablehlo.add %35173, %35740 : tensor<1x1x2048x2048xbf16> loc(#loc4155)
      %35742 = arith.select %ptrue, %35740, %35740 : tensor<1x1x2048x2048xbf16>
      %35743:2 = func.call @_einsum_859(%30575#0, %34420, %35742) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4156)
      %35744 = func.call @_where_860(%30575#1, %35743#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %35745 = arith.select %ptrue, %35743#0, %35743#0 : tensor<1x1x2048x32x64xbf16>
      %35746:2 = func.call @_einsum_861(%34418, %34375, %35745) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4157)
      %35747 = arith.select %ptrue, %35746#1, %35746#1 : tensor<1x1x2048x32x64xbf16>
      %35748 = stablehlo.broadcast_in_dim %35747, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4158)
      %35749 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35750 = stablehlo.pad %35748, %35749, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4159)
      %35751 = stablehlo.convert %35746#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4160)
      %35752 = stablehlo.broadcast_in_dim %34417, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3641)
      %35753 = stablehlo.multiply %35751, %35752 : tensor<1x1x32x2048x2048xf32> loc(#loc3641)
      %35754 = stablehlo.multiply %35753, %34408 : tensor<1x1x32x2048x2048xf32> loc(#loc3641)
      %35755 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3639)
      %35756 = stablehlo.reduce(%35754 init: %35755) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3639)
        stablehlo.return %39096 : tensor<f32> loc(#loc3639)
      } loc(#loc3639)
      %35757 = stablehlo.reshape %35756 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4161)
      %35758 = stablehlo.negate %35757 : tensor<1x1x32x2048x1xf32> loc(#loc4162)
      %35759 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4163)
      %35760 = stablehlo.reduce(%35758 init: %35759) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4163)
        stablehlo.return %39096 : tensor<f32> loc(#loc4163)
      } loc(#loc4163)
      %35761 = stablehlo.broadcast_in_dim %35760, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4164)
      %35762 = stablehlo.broadcast_in_dim %35761, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4165)
      %35763 = stablehlo.broadcast_in_dim %34411, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3640)
      %35764 = stablehlo.divide %35751, %35763 : tensor<1x1x32x2048x2048xf32> loc(#loc3640)
      %35765 = stablehlo.add %35762, %35764 : tensor<1x1x32x2048x2048xf32> loc(#loc4166)
      %35766 = stablehlo.multiply %35765, %34408 : tensor<1x1x32x2048x2048xf32> loc(#loc3641)
      %35767 = func.call @_where_862(%34402#1, %35766) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4167)
      %35768 = stablehlo.convert %35767 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4168)
      %35769 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %35770 = stablehlo.broadcast_in_dim %35769, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3629)
      %35771 = stablehlo.multiply %35770, %35768 : tensor<1x1x32x2048x2048xbf16> loc(#loc3629)
      %35772 = stablehlo.multiply %35771, %34390 : tensor<1x1x32x2048x2048xbf16> loc(#loc3629)
      %35773 = stablehlo.multiply %35772, %34387 : tensor<1x1x32x2048x2048xbf16> loc(#loc3629)
      %35774 = stablehlo.add %35772, %35773 : tensor<1x1x32x2048x2048xbf16> loc(#loc4169)
      %35775 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %35776 = stablehlo.broadcast_in_dim %35775, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3626)
      %35777 = stablehlo.divide %35774, %35776 : tensor<1x1x32x2048x2048xbf16> loc(#loc3626)
      %35778:2 = func.call @_einsum_863(%34374, %34382, %35777) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4170)
      %35779 = arith.select %ptrue, %35778#1, %35778#1 : tensor<1x1x2048x32x64xbf16>
      %35780 = stablehlo.broadcast_in_dim %35779, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4158)
      %35781 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35782 = stablehlo.pad %35780, %35781, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4171)
      %35783 = stablehlo.add %35750, %35782 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4172)
      %35784 = stablehlo.multiply %34373, %35778#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3624)
      %35785 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4173)
      %35786 = stablehlo.reduce(%35784 init: %35785) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4173)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4173)
      } loc(#loc4173)
      %35787 = stablehlo.reshape %35786 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4174)
      %35788 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4175)
      %35789 = stablehlo.reduce(%35787 init: %35788) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4175)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4175)
      } loc(#loc4175)
      %35790 = stablehlo.broadcast_in_dim %35789, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4176)
      %35791 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %35792 = stablehlo.broadcast_in_dim %35791, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3622)
      %35793 = stablehlo.multiply %35792, %35790 : tensor<1x64xbf16> loc(#loc3622)
      %35794 = func.call @softplus_864(%34376#1, %35793) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4177)
      %35795 = func.call @_where_865(%30574#1, %35794) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %35796 = stablehlo.broadcast_in_dim %34380, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3624)
      %35797 = stablehlo.multiply %35778#0, %35796 : tensor<1x1x2048x32x64xbf16> loc(#loc3624)
      %35798 = arith.select %ptrue, %35797, %35797 : tensor<1x1x2048x32x64xbf16>
      %35799 = stablehlo.broadcast_in_dim %35798, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4158)
      %35800 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35801 = stablehlo.pad %35799, %35800, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4178)
      %35802 = stablehlo.add %35783, %35801 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4172)
      %35803:2 = func.call @_einsum_866(%34365, %30573#0, %35802) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4179)
      %35804 = func.call @_where_867(%30573#1, %35803#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %35805 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4180)
      %35806 = stablehlo.reduce(%35803#0 init: %35805) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4180)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4180)
      } loc(#loc4180)
      %35807 = stablehlo.reshape %35806 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4181)
      %35808 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4182)
      %35809 = stablehlo.reduce(%35807 init: %35808) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4182)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4182)
      } loc(#loc4182)
      %35810 = stablehlo.broadcast_in_dim %35809, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4183)
      %35811 = func.call @_where_868(%30571#1, %35810) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35812 = stablehlo.multiply %34356, %35803#0 : tensor<1x1x2048x2048xbf16> loc(#loc3612)
      %35813 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4184)
      %35814 = stablehlo.reduce(%35812 init: %35813) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4184)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4184)
      } loc(#loc4184)
      %35815 = stablehlo.reshape %35814 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4185)
      %35816 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4186)
      %35817 = stablehlo.reduce(%35815 init: %35816) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4186)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4186)
      } loc(#loc4186)
      %35818 = stablehlo.broadcast_in_dim %35817, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4187)
      %35819 = func.call @_where_869(%30572#1, %35818) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35820 = stablehlo.broadcast_in_dim %34360, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3612)
      %35821 = stablehlo.multiply %35803#0, %35820 : tensor<1x1x2048x2048xbf16> loc(#loc3612)
      %35822 = stablehlo.multiply %34346, %35821 : tensor<1x1x2048x2048xbf16> loc(#loc3609)
      %35823 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4188)
      %35824 = stablehlo.reduce(%35822 init: %35823) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4188)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4188)
      } loc(#loc4188)
      %35825 = stablehlo.reshape %35824 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4189)
      %35826 = stablehlo.multiply %35825, %34354 : tensor<1x1x2048x1xbf16> loc(#loc3609)
      %35827 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35828 = stablehlo.broadcast_in_dim %35827, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3603)
      %35829 = stablehlo.divide %35826, %35828 : tensor<1x1x2048x1xbf16> loc(#loc3603)
      %35830 = stablehlo.convert %35829 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3599)
      %35831 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4190)
      %35832 = stablehlo.reduce(%35830 init: %35831) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4190)
        stablehlo.return %39096 : tensor<f32> loc(#loc4190)
      } loc(#loc4190)
      %35833 = stablehlo.broadcast_in_dim %35832, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4191)
      %35834 = stablehlo.broadcast_in_dim %35833, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4192)
      %35835 = stablehlo.convert %35834 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3602)
      %35836 = stablehlo.multiply %35835, %34336 : tensor<1x1x2048x2048xbf16> loc(#loc3605)
      %35837 = stablehlo.negate %35836 : tensor<1x1x2048x2048xbf16> loc(#loc4193)
      %35838 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4194)
      %35839 = stablehlo.reduce(%35837 init: %35838) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4194)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4194)
      } loc(#loc4194)
      %35840 = stablehlo.reshape %35839 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4195)
      %35841 = stablehlo.add %35741, %35836 : tensor<1x1x2048x2048xbf16> loc(#loc4196)
      %35842 = stablehlo.broadcast_in_dim %34350, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3609)
      %35843 = stablehlo.multiply %35821, %35842 : tensor<1x1x2048x2048xbf16> loc(#loc3609)
      %35844 = stablehlo.negate %35843 : tensor<1x1x2048x2048xbf16> loc(#loc4193)
      %35845 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4194)
      %35846 = stablehlo.reduce(%35844 init: %35845) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4194)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4194)
      } loc(#loc4194)
      %35847 = stablehlo.reshape %35846 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4195)
      %35848 = stablehlo.add %35841, %35843 : tensor<1x1x2048x2048xbf16> loc(#loc4196)
      %35849 = stablehlo.add %35840, %35847 : tensor<1x1x2048x1xbf16> loc(#loc4196)
      %35850 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35851 = stablehlo.broadcast_in_dim %35850, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3603)
      %35852 = stablehlo.divide %35849, %35851 : tensor<1x1x2048x1xbf16> loc(#loc3603)
      %35853 = stablehlo.convert %35852 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3599)
      %35854 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4190)
      %35855 = stablehlo.reduce(%35853 init: %35854) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4190)
        stablehlo.return %39096 : tensor<f32> loc(#loc4190)
      } loc(#loc4190)
      %35856 = stablehlo.broadcast_in_dim %35855, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4191)
      %35857 = stablehlo.broadcast_in_dim %35856, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4192)
      %35858 = stablehlo.convert %35857 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3602)
      %35859 = stablehlo.add %35848, %35858 : tensor<1x1x2048x2048xbf16> loc(#loc4197)
      %35860 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %35861 = stablehlo.broadcast_in_dim %35860, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3574)
      %35862 = stablehlo.multiply %35859, %35861 : tensor<1x1x2048x2048xbf16> loc(#loc3574)
      %35863 = stablehlo.broadcast_in_dim %34258, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3573)
      %35864 = stablehlo.multiply %35862, %35863 : tensor<1x1x2048x2048xbf16> loc(#loc3573)
      %35865 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4198)
      %35866 = stablehlo.reduce(%35864 init: %35865) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4198)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4198)
      } loc(#loc4198)
      %35867 = stablehlo.reshape %35866 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4199)
      %35868 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4200)
      %35869 = stablehlo.reduce(%35867 init: %35868) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4200)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4200)
      } loc(#loc4200)
      %35870 = stablehlo.broadcast_in_dim %35869, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4201)
      %35871 = func.call @_where_870(%30556#1, %35870) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35872 = arith.select %ptrue, %35864, %35864 : tensor<1x1x2048x2048xbf16>
      %35873:2 = func.call @_einsum_871(%30557#0, %34248, %35872) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4202)
      %35874 = func.call @_where_872(%30557#1, %35873#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %35875 = stablehlo.broadcast_in_dim %34246, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3567)
      %35876 = stablehlo.multiply %35873#0, %35875 : tensor<1x1x2048x8192xbf16> loc(#loc3567)
      %35877 = stablehlo.multiply %34217, %35876 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35878 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %35879 = stablehlo.broadcast_in_dim %35878, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35880 = stablehlo.multiply %35879, %35877 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35881 = stablehlo.multiply %35880, %34234 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35882 = stablehlo.multiply %35881, %34231 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35883 = stablehlo.add %35881, %35882 : tensor<1x1x2048x8192xbf16> loc(#loc4203)
      %35884 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %35885 = stablehlo.broadcast_in_dim %35884, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35886 = stablehlo.multiply %35885, %35883 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35887 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %35888 = stablehlo.broadcast_in_dim %35887, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35889 = stablehlo.multiply %35888, %35886 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35890 = stablehlo.multiply %35889, %34223 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35891 = stablehlo.add %35886, %35890 : tensor<1x1x2048x8192xbf16> loc(#loc4203)
      %35892 = stablehlo.multiply %35876, %34240 : tensor<1x1x2048x8192xbf16> loc(#loc3560)
      %35893 = stablehlo.add %35891, %35892 : tensor<1x1x2048x8192xbf16> loc(#loc4203)
      %35894 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4204)
      %35895 = stablehlo.reduce(%35893 init: %35894) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4204)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4204)
      } loc(#loc4204)
      %35896 = stablehlo.reshape %35895 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4205)
      %35897 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4206)
      %35898 = stablehlo.reduce(%35896 init: %35897) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4206)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4206)
      } loc(#loc4206)
      %35899 = stablehlo.broadcast_in_dim %35898, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4207)
      %35900 = func.call @_where_873(%30554#1, %35899) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %35901 = arith.select %ptrue, %35893, %35893 : tensor<1x1x2048x8192xbf16>
      %35902:2 = func.call @_einsum_874(%30555#0, %34212, %35901) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4208)
      %35903 = func.call @_where_875(%30555#1, %35902#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %35904 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4209)
      %35905 = stablehlo.reduce(%35902#0 init: %35904) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4209)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4209)
      } loc(#loc4209)
      %35906 = stablehlo.reshape %35905 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4210)
      %35907 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4211)
      %35908 = stablehlo.reduce(%35906 init: %35907) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4211)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4211)
      } loc(#loc4211)
      %35909 = stablehlo.broadcast_in_dim %35908, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4212)
      %35910 = func.call @_where_876(%30558#1, %35909) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35911 = stablehlo.multiply %34203, %35902#0 : tensor<1x1x2048x2048xbf16> loc(#loc3553)
      %35912 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4213)
      %35913 = stablehlo.reduce(%35911 init: %35912) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4213)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4213)
      } loc(#loc4213)
      %35914 = stablehlo.reshape %35913 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4214)
      %35915 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4215)
      %35916 = stablehlo.reduce(%35914 init: %35915) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4215)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4215)
      } loc(#loc4215)
      %35917 = stablehlo.broadcast_in_dim %35916, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4216)
      %35918 = func.call @_where_877(%30559#1, %35917) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %35919 = stablehlo.broadcast_in_dim %34207, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3553)
      %35920 = stablehlo.multiply %35902#0, %35919 : tensor<1x1x2048x2048xbf16> loc(#loc3553)
      %35921 = stablehlo.multiply %34193, %35920 : tensor<1x1x2048x2048xbf16> loc(#loc3550)
      %35922 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4217)
      %35923 = stablehlo.reduce(%35921 init: %35922) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4217)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4217)
      } loc(#loc4217)
      %35924 = stablehlo.reshape %35923 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4218)
      %35925 = stablehlo.multiply %35924, %34201 : tensor<1x1x2048x1xbf16> loc(#loc3550)
      %35926 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35927 = stablehlo.broadcast_in_dim %35926, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3544)
      %35928 = stablehlo.divide %35925, %35927 : tensor<1x1x2048x1xbf16> loc(#loc3544)
      %35929 = stablehlo.convert %35928 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3540)
      %35930 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4219)
      %35931 = stablehlo.reduce(%35929 init: %35930) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4219)
        stablehlo.return %39096 : tensor<f32> loc(#loc4219)
      } loc(#loc4219)
      %35932 = stablehlo.broadcast_in_dim %35931, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4220)
      %35933 = stablehlo.broadcast_in_dim %35932, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4221)
      %35934 = stablehlo.convert %35933 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3543)
      %35935 = stablehlo.multiply %35934, %34183 : tensor<1x1x2048x2048xbf16> loc(#loc3546)
      %35936 = stablehlo.negate %35935 : tensor<1x1x2048x2048xbf16> loc(#loc4222)
      %35937 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4223)
      %35938 = stablehlo.reduce(%35936 init: %35937) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4223)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4223)
      } loc(#loc4223)
      %35939 = stablehlo.reshape %35938 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4224)
      %35940 = stablehlo.add %35859, %35935 : tensor<1x1x2048x2048xbf16> loc(#loc4225)
      %35941 = stablehlo.broadcast_in_dim %34197, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3550)
      %35942 = stablehlo.multiply %35920, %35941 : tensor<1x1x2048x2048xbf16> loc(#loc3550)
      %35943 = stablehlo.negate %35942 : tensor<1x1x2048x2048xbf16> loc(#loc4222)
      %35944 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4223)
      %35945 = stablehlo.reduce(%35943 init: %35944) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4223)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4223)
      } loc(#loc4223)
      %35946 = stablehlo.reshape %35945 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4224)
      %35947 = stablehlo.add %35940, %35942 : tensor<1x1x2048x2048xbf16> loc(#loc4225)
      %35948 = stablehlo.add %35939, %35946 : tensor<1x1x2048x1xbf16> loc(#loc4225)
      %35949 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %35950 = stablehlo.broadcast_in_dim %35949, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3544)
      %35951 = stablehlo.divide %35948, %35950 : tensor<1x1x2048x1xbf16> loc(#loc3544)
      %35952 = stablehlo.convert %35951 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3540)
      %35953 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4219)
      %35954 = stablehlo.reduce(%35952 init: %35953) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4219)
        stablehlo.return %39096 : tensor<f32> loc(#loc4219)
      } loc(#loc4219)
      %35955 = stablehlo.broadcast_in_dim %35954, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4220)
      %35956 = stablehlo.broadcast_in_dim %35955, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4221)
      %35957 = stablehlo.convert %35956 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3543)
      %35958 = stablehlo.add %35947, %35957 : tensor<1x1x2048x2048xbf16> loc(#loc4226)
      %35959 = stablehlo.add %35214, %35958 : tensor<1x1x2048x2048xbf16> loc(#loc4227)
      %35960 = arith.select %ptrue, %35958, %35958 : tensor<1x1x2048x2048xbf16>
      %35961:2 = func.call @_einsum_878(%30564#0, %34165, %35960) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4228)
      %35962 = func.call @_where_879(%30564#1, %35961#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %35963 = arith.select %ptrue, %35961#0, %35961#0 : tensor<1x1x2048x32x64xbf16>
      %35964:2 = func.call @_einsum_880(%34163, %34120, %35963) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4229)
      %35965 = arith.select %ptrue, %35964#1, %35964#1 : tensor<1x1x2048x32x64xbf16>
      %35966 = stablehlo.broadcast_in_dim %35965, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4230)
      %35967 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %35968 = stablehlo.pad %35966, %35967, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4231)
      %35969 = stablehlo.convert %35964#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4232)
      %35970 = stablehlo.broadcast_in_dim %34162, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3531)
      %35971 = stablehlo.multiply %35969, %35970 : tensor<1x1x32x2048x2048xf32> loc(#loc3531)
      %35972 = stablehlo.multiply %35971, %34153 : tensor<1x1x32x2048x2048xf32> loc(#loc3531)
      %35973 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3529)
      %35974 = stablehlo.reduce(%35972 init: %35973) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3529)
        stablehlo.return %39096 : tensor<f32> loc(#loc3529)
      } loc(#loc3529)
      %35975 = stablehlo.reshape %35974 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4233)
      %35976 = stablehlo.negate %35975 : tensor<1x1x32x2048x1xf32> loc(#loc4234)
      %35977 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4235)
      %35978 = stablehlo.reduce(%35976 init: %35977) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4235)
        stablehlo.return %39096 : tensor<f32> loc(#loc4235)
      } loc(#loc4235)
      %35979 = stablehlo.broadcast_in_dim %35978, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4236)
      %35980 = stablehlo.broadcast_in_dim %35979, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4237)
      %35981 = stablehlo.broadcast_in_dim %34156, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3530)
      %35982 = stablehlo.divide %35969, %35981 : tensor<1x1x32x2048x2048xf32> loc(#loc3530)
      %35983 = stablehlo.add %35980, %35982 : tensor<1x1x32x2048x2048xf32> loc(#loc4238)
      %35984 = stablehlo.multiply %35983, %34153 : tensor<1x1x32x2048x2048xf32> loc(#loc3531)
      %35985 = func.call @_where_881(%34147#1, %35984) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4239)
      %35986 = stablehlo.convert %35985 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4240)
      %35987 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %35988 = stablehlo.broadcast_in_dim %35987, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3519)
      %35989 = stablehlo.multiply %35988, %35986 : tensor<1x1x32x2048x2048xbf16> loc(#loc3519)
      %35990 = stablehlo.multiply %35989, %34135 : tensor<1x1x32x2048x2048xbf16> loc(#loc3519)
      %35991 = stablehlo.multiply %35990, %34132 : tensor<1x1x32x2048x2048xbf16> loc(#loc3519)
      %35992 = stablehlo.add %35990, %35991 : tensor<1x1x32x2048x2048xbf16> loc(#loc4241)
      %35993 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %35994 = stablehlo.broadcast_in_dim %35993, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3516)
      %35995 = stablehlo.divide %35992, %35994 : tensor<1x1x32x2048x2048xbf16> loc(#loc3516)
      %35996:2 = func.call @_einsum_882(%34119, %34127, %35995) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4242)
      %35997 = arith.select %ptrue, %35996#1, %35996#1 : tensor<1x1x2048x32x64xbf16>
      %35998 = stablehlo.broadcast_in_dim %35997, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4230)
      %35999 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36000 = stablehlo.pad %35998, %35999, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4243)
      %36001 = stablehlo.add %35968, %36000 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4244)
      %36002 = stablehlo.multiply %34118, %35996#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3514)
      %36003 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4245)
      %36004 = stablehlo.reduce(%36002 init: %36003) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4245)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4245)
      } loc(#loc4245)
      %36005 = stablehlo.reshape %36004 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4246)
      %36006 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4247)
      %36007 = stablehlo.reduce(%36005 init: %36006) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4247)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4247)
      } loc(#loc4247)
      %36008 = stablehlo.broadcast_in_dim %36007, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4248)
      %36009 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %36010 = stablehlo.broadcast_in_dim %36009, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3512)
      %36011 = stablehlo.multiply %36010, %36008 : tensor<1x64xbf16> loc(#loc3512)
      %36012 = func.call @softplus_883(%34121#1, %36011) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4249)
      %36013 = func.call @_where_884(%30563#1, %36012) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %36014 = stablehlo.broadcast_in_dim %34125, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3514)
      %36015 = stablehlo.multiply %35996#0, %36014 : tensor<1x1x2048x32x64xbf16> loc(#loc3514)
      %36016 = arith.select %ptrue, %36015, %36015 : tensor<1x1x2048x32x64xbf16>
      %36017 = stablehlo.broadcast_in_dim %36016, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4230)
      %36018 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36019 = stablehlo.pad %36017, %36018, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4250)
      %36020 = stablehlo.add %36001, %36019 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4244)
      %36021:2 = func.call @_einsum_885(%34110, %30562#0, %36020) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4251)
      %36022 = func.call @_where_886(%30562#1, %36021#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %36023 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4252)
      %36024 = stablehlo.reduce(%36021#0 init: %36023) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4252)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4252)
      } loc(#loc4252)
      %36025 = stablehlo.reshape %36024 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4253)
      %36026 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4254)
      %36027 = stablehlo.reduce(%36025 init: %36026) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4254)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4254)
      } loc(#loc4254)
      %36028 = stablehlo.broadcast_in_dim %36027, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4255)
      %36029 = func.call @_where_887(%30560#1, %36028) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36030 = stablehlo.multiply %34101, %36021#0 : tensor<1x1x2048x2048xbf16> loc(#loc3502)
      %36031 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4256)
      %36032 = stablehlo.reduce(%36030 init: %36031) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4256)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4256)
      } loc(#loc4256)
      %36033 = stablehlo.reshape %36032 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4257)
      %36034 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4258)
      %36035 = stablehlo.reduce(%36033 init: %36034) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4258)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4258)
      } loc(#loc4258)
      %36036 = stablehlo.broadcast_in_dim %36035, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4259)
      %36037 = func.call @_where_888(%30561#1, %36036) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36038 = stablehlo.broadcast_in_dim %34105, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3502)
      %36039 = stablehlo.multiply %36021#0, %36038 : tensor<1x1x2048x2048xbf16> loc(#loc3502)
      %36040 = stablehlo.multiply %34091, %36039 : tensor<1x1x2048x2048xbf16> loc(#loc3499)
      %36041 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4260)
      %36042 = stablehlo.reduce(%36040 init: %36041) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4260)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4260)
      } loc(#loc4260)
      %36043 = stablehlo.reshape %36042 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4261)
      %36044 = stablehlo.multiply %36043, %34099 : tensor<1x1x2048x1xbf16> loc(#loc3499)
      %36045 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36046 = stablehlo.broadcast_in_dim %36045, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3493)
      %36047 = stablehlo.divide %36044, %36046 : tensor<1x1x2048x1xbf16> loc(#loc3493)
      %36048 = stablehlo.convert %36047 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3489)
      %36049 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4262)
      %36050 = stablehlo.reduce(%36048 init: %36049) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4262)
        stablehlo.return %39096 : tensor<f32> loc(#loc4262)
      } loc(#loc4262)
      %36051 = stablehlo.broadcast_in_dim %36050, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4263)
      %36052 = stablehlo.broadcast_in_dim %36051, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4264)
      %36053 = stablehlo.convert %36052 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3492)
      %36054 = stablehlo.multiply %36053, %34081 : tensor<1x1x2048x2048xbf16> loc(#loc3495)
      %36055 = stablehlo.negate %36054 : tensor<1x1x2048x2048xbf16> loc(#loc4265)
      %36056 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4266)
      %36057 = stablehlo.reduce(%36055 init: %36056) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4266)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4266)
      } loc(#loc4266)
      %36058 = stablehlo.reshape %36057 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4267)
      %36059 = stablehlo.add %35959, %36054 : tensor<1x1x2048x2048xbf16> loc(#loc4268)
      %36060 = stablehlo.broadcast_in_dim %34095, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3499)
      %36061 = stablehlo.multiply %36039, %36060 : tensor<1x1x2048x2048xbf16> loc(#loc3499)
      %36062 = stablehlo.negate %36061 : tensor<1x1x2048x2048xbf16> loc(#loc4265)
      %36063 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4266)
      %36064 = stablehlo.reduce(%36062 init: %36063) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4266)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4266)
      } loc(#loc4266)
      %36065 = stablehlo.reshape %36064 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4267)
      %36066 = stablehlo.add %36059, %36061 : tensor<1x1x2048x2048xbf16> loc(#loc4268)
      %36067 = stablehlo.add %36058, %36065 : tensor<1x1x2048x1xbf16> loc(#loc4268)
      %36068 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36069 = stablehlo.broadcast_in_dim %36068, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3493)
      %36070 = stablehlo.divide %36067, %36069 : tensor<1x1x2048x1xbf16> loc(#loc3493)
      %36071 = stablehlo.convert %36070 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3489)
      %36072 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4262)
      %36073 = stablehlo.reduce(%36071 init: %36072) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4262)
        stablehlo.return %39096 : tensor<f32> loc(#loc4262)
      } loc(#loc4262)
      %36074 = stablehlo.broadcast_in_dim %36073, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4263)
      %36075 = stablehlo.broadcast_in_dim %36074, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4264)
      %36076 = stablehlo.convert %36075 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3492)
      %36077 = stablehlo.add %36066, %36076 : tensor<1x1x2048x2048xbf16> loc(#loc4269)
      %36078 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %36079 = stablehlo.broadcast_in_dim %36078, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3464)
      %36080 = stablehlo.multiply %36077, %36079 : tensor<1x1x2048x2048xbf16> loc(#loc3464)
      %36081 = stablehlo.broadcast_in_dim %34003, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3463)
      %36082 = stablehlo.multiply %36080, %36081 : tensor<1x1x2048x2048xbf16> loc(#loc3463)
      %36083 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4270)
      %36084 = stablehlo.reduce(%36082 init: %36083) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4270)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4270)
      } loc(#loc4270)
      %36085 = stablehlo.reshape %36084 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4271)
      %36086 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4272)
      %36087 = stablehlo.reduce(%36085 init: %36086) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4272)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4272)
      } loc(#loc4272)
      %36088 = stablehlo.broadcast_in_dim %36087, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4273)
      %36089 = func.call @_where_889(%30545#1, %36088) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36090 = arith.select %ptrue, %36082, %36082 : tensor<1x1x2048x2048xbf16>
      %36091:2 = func.call @_einsum_890(%30546#0, %33993, %36090) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4274)
      %36092 = func.call @_where_891(%30546#1, %36091#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %36093 = stablehlo.broadcast_in_dim %33991, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3457)
      %36094 = stablehlo.multiply %36091#0, %36093 : tensor<1x1x2048x8192xbf16> loc(#loc3457)
      %36095 = stablehlo.multiply %33962, %36094 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36096 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %36097 = stablehlo.broadcast_in_dim %36096, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36098 = stablehlo.multiply %36097, %36095 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36099 = stablehlo.multiply %36098, %33979 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36100 = stablehlo.multiply %36099, %33976 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36101 = stablehlo.add %36099, %36100 : tensor<1x1x2048x8192xbf16> loc(#loc4275)
      %36102 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %36103 = stablehlo.broadcast_in_dim %36102, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36104 = stablehlo.multiply %36103, %36101 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36105 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %36106 = stablehlo.broadcast_in_dim %36105, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36107 = stablehlo.multiply %36106, %36104 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36108 = stablehlo.multiply %36107, %33968 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36109 = stablehlo.add %36104, %36108 : tensor<1x1x2048x8192xbf16> loc(#loc4275)
      %36110 = stablehlo.multiply %36094, %33985 : tensor<1x1x2048x8192xbf16> loc(#loc3450)
      %36111 = stablehlo.add %36109, %36110 : tensor<1x1x2048x8192xbf16> loc(#loc4275)
      %36112 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4276)
      %36113 = stablehlo.reduce(%36111 init: %36112) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4276)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4276)
      } loc(#loc4276)
      %36114 = stablehlo.reshape %36113 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4277)
      %36115 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4278)
      %36116 = stablehlo.reduce(%36114 init: %36115) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4278)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4278)
      } loc(#loc4278)
      %36117 = stablehlo.broadcast_in_dim %36116, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4279)
      %36118 = func.call @_where_892(%30543#1, %36117) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %36119 = arith.select %ptrue, %36111, %36111 : tensor<1x1x2048x8192xbf16>
      %36120:2 = func.call @_einsum_893(%30544#0, %33957, %36119) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4280)
      %36121 = func.call @_where_894(%30544#1, %36120#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %36122 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4281)
      %36123 = stablehlo.reduce(%36120#0 init: %36122) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4281)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4281)
      } loc(#loc4281)
      %36124 = stablehlo.reshape %36123 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4282)
      %36125 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4283)
      %36126 = stablehlo.reduce(%36124 init: %36125) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4283)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4283)
      } loc(#loc4283)
      %36127 = stablehlo.broadcast_in_dim %36126, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4284)
      %36128 = func.call @_where_895(%30547#1, %36127) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36129 = stablehlo.multiply %33948, %36120#0 : tensor<1x1x2048x2048xbf16> loc(#loc3443)
      %36130 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4285)
      %36131 = stablehlo.reduce(%36129 init: %36130) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4285)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4285)
      } loc(#loc4285)
      %36132 = stablehlo.reshape %36131 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4286)
      %36133 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4287)
      %36134 = stablehlo.reduce(%36132 init: %36133) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4287)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4287)
      } loc(#loc4287)
      %36135 = stablehlo.broadcast_in_dim %36134, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4288)
      %36136 = func.call @_where_896(%30548#1, %36135) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36137 = stablehlo.broadcast_in_dim %33952, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3443)
      %36138 = stablehlo.multiply %36120#0, %36137 : tensor<1x1x2048x2048xbf16> loc(#loc3443)
      %36139 = stablehlo.multiply %33938, %36138 : tensor<1x1x2048x2048xbf16> loc(#loc3440)
      %36140 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4289)
      %36141 = stablehlo.reduce(%36139 init: %36140) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4289)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4289)
      } loc(#loc4289)
      %36142 = stablehlo.reshape %36141 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4290)
      %36143 = stablehlo.multiply %36142, %33946 : tensor<1x1x2048x1xbf16> loc(#loc3440)
      %36144 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36145 = stablehlo.broadcast_in_dim %36144, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3434)
      %36146 = stablehlo.divide %36143, %36145 : tensor<1x1x2048x1xbf16> loc(#loc3434)
      %36147 = stablehlo.convert %36146 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3430)
      %36148 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4291)
      %36149 = stablehlo.reduce(%36147 init: %36148) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4291)
        stablehlo.return %39096 : tensor<f32> loc(#loc4291)
      } loc(#loc4291)
      %36150 = stablehlo.broadcast_in_dim %36149, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4292)
      %36151 = stablehlo.broadcast_in_dim %36150, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4293)
      %36152 = stablehlo.convert %36151 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3433)
      %36153 = stablehlo.multiply %36152, %33928 : tensor<1x1x2048x2048xbf16> loc(#loc3436)
      %36154 = stablehlo.negate %36153 : tensor<1x1x2048x2048xbf16> loc(#loc4294)
      %36155 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4295)
      %36156 = stablehlo.reduce(%36154 init: %36155) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4295)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4295)
      } loc(#loc4295)
      %36157 = stablehlo.reshape %36156 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4296)
      %36158 = stablehlo.add %36077, %36153 : tensor<1x1x2048x2048xbf16> loc(#loc4297)
      %36159 = stablehlo.broadcast_in_dim %33942, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3440)
      %36160 = stablehlo.multiply %36138, %36159 : tensor<1x1x2048x2048xbf16> loc(#loc3440)
      %36161 = stablehlo.negate %36160 : tensor<1x1x2048x2048xbf16> loc(#loc4294)
      %36162 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4295)
      %36163 = stablehlo.reduce(%36161 init: %36162) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4295)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4295)
      } loc(#loc4295)
      %36164 = stablehlo.reshape %36163 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4296)
      %36165 = stablehlo.add %36158, %36160 : tensor<1x1x2048x2048xbf16> loc(#loc4297)
      %36166 = stablehlo.add %36157, %36164 : tensor<1x1x2048x1xbf16> loc(#loc4297)
      %36167 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36168 = stablehlo.broadcast_in_dim %36167, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3434)
      %36169 = stablehlo.divide %36166, %36168 : tensor<1x1x2048x1xbf16> loc(#loc3434)
      %36170 = stablehlo.convert %36169 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3430)
      %36171 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4291)
      %36172 = stablehlo.reduce(%36170 init: %36171) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4291)
        stablehlo.return %39096 : tensor<f32> loc(#loc4291)
      } loc(#loc4291)
      %36173 = stablehlo.broadcast_in_dim %36172, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4292)
      %36174 = stablehlo.broadcast_in_dim %36173, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4293)
      %36175 = stablehlo.convert %36174 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3433)
      %36176 = stablehlo.add %36165, %36175 : tensor<1x1x2048x2048xbf16> loc(#loc4298)
      %36177 = stablehlo.add %35255, %36176 : tensor<1x1x2048x2048xbf16> loc(#loc4299)
      %36178 = arith.select %ptrue, %36176, %36176 : tensor<1x1x2048x2048xbf16>
      %36179:2 = func.call @_einsum_897(%30553#0, %33910, %36178) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4300)
      %36180 = func.call @_where_898(%30553#1, %36179#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %36181 = arith.select %ptrue, %36179#0, %36179#0 : tensor<1x1x2048x32x64xbf16>
      %36182:2 = func.call @_einsum_899(%33908, %33865, %36181) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4301)
      %36183 = arith.select %ptrue, %36182#1, %36182#1 : tensor<1x1x2048x32x64xbf16>
      %36184 = stablehlo.broadcast_in_dim %36183, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4302)
      %36185 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36186 = stablehlo.pad %36184, %36185, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4303)
      %36187 = stablehlo.convert %36182#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4304)
      %36188 = stablehlo.broadcast_in_dim %33907, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3421)
      %36189 = stablehlo.multiply %36187, %36188 : tensor<1x1x32x2048x2048xf32> loc(#loc3421)
      %36190 = stablehlo.multiply %36189, %33898 : tensor<1x1x32x2048x2048xf32> loc(#loc3421)
      %36191 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3419)
      %36192 = stablehlo.reduce(%36190 init: %36191) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3419)
        stablehlo.return %39096 : tensor<f32> loc(#loc3419)
      } loc(#loc3419)
      %36193 = stablehlo.reshape %36192 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4305)
      %36194 = stablehlo.negate %36193 : tensor<1x1x32x2048x1xf32> loc(#loc4306)
      %36195 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4307)
      %36196 = stablehlo.reduce(%36194 init: %36195) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4307)
        stablehlo.return %39096 : tensor<f32> loc(#loc4307)
      } loc(#loc4307)
      %36197 = stablehlo.broadcast_in_dim %36196, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4308)
      %36198 = stablehlo.broadcast_in_dim %36197, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4309)
      %36199 = stablehlo.broadcast_in_dim %33901, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3420)
      %36200 = stablehlo.divide %36187, %36199 : tensor<1x1x32x2048x2048xf32> loc(#loc3420)
      %36201 = stablehlo.add %36198, %36200 : tensor<1x1x32x2048x2048xf32> loc(#loc4310)
      %36202 = stablehlo.multiply %36201, %33898 : tensor<1x1x32x2048x2048xf32> loc(#loc3421)
      %36203 = func.call @_where_900(%33892#1, %36202) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4311)
      %36204 = stablehlo.convert %36203 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4312)
      %36205 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36206 = stablehlo.broadcast_in_dim %36205, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3409)
      %36207 = stablehlo.multiply %36206, %36204 : tensor<1x1x32x2048x2048xbf16> loc(#loc3409)
      %36208 = stablehlo.multiply %36207, %33880 : tensor<1x1x32x2048x2048xbf16> loc(#loc3409)
      %36209 = stablehlo.multiply %36208, %33877 : tensor<1x1x32x2048x2048xbf16> loc(#loc3409)
      %36210 = stablehlo.add %36208, %36209 : tensor<1x1x32x2048x2048xbf16> loc(#loc4313)
      %36211 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36212 = stablehlo.broadcast_in_dim %36211, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3406)
      %36213 = stablehlo.divide %36210, %36212 : tensor<1x1x32x2048x2048xbf16> loc(#loc3406)
      %36214:2 = func.call @_einsum_901(%33864, %33872, %36213) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4314)
      %36215 = arith.select %ptrue, %36214#1, %36214#1 : tensor<1x1x2048x32x64xbf16>
      %36216 = stablehlo.broadcast_in_dim %36215, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4302)
      %36217 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36218 = stablehlo.pad %36216, %36217, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4315)
      %36219 = stablehlo.add %36186, %36218 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4316)
      %36220 = stablehlo.multiply %33863, %36214#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3404)
      %36221 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4317)
      %36222 = stablehlo.reduce(%36220 init: %36221) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4317)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4317)
      } loc(#loc4317)
      %36223 = stablehlo.reshape %36222 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4318)
      %36224 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4319)
      %36225 = stablehlo.reduce(%36223 init: %36224) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4319)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4319)
      } loc(#loc4319)
      %36226 = stablehlo.broadcast_in_dim %36225, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4320)
      %36227 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %36228 = stablehlo.broadcast_in_dim %36227, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3402)
      %36229 = stablehlo.multiply %36228, %36226 : tensor<1x64xbf16> loc(#loc3402)
      %36230 = func.call @softplus_902(%33866#1, %36229) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4321)
      %36231 = func.call @_where_903(%30552#1, %36230) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %36232 = stablehlo.broadcast_in_dim %33870, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3404)
      %36233 = stablehlo.multiply %36214#0, %36232 : tensor<1x1x2048x32x64xbf16> loc(#loc3404)
      %36234 = arith.select %ptrue, %36233, %36233 : tensor<1x1x2048x32x64xbf16>
      %36235 = stablehlo.broadcast_in_dim %36234, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4302)
      %36236 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36237 = stablehlo.pad %36235, %36236, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4322)
      %36238 = stablehlo.add %36219, %36237 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4316)
      %36239:2 = func.call @_einsum_904(%33855, %30551#0, %36238) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4323)
      %36240 = func.call @_where_905(%30551#1, %36239#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %36241 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4324)
      %36242 = stablehlo.reduce(%36239#0 init: %36241) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4324)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4324)
      } loc(#loc4324)
      %36243 = stablehlo.reshape %36242 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4325)
      %36244 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4326)
      %36245 = stablehlo.reduce(%36243 init: %36244) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4326)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4326)
      } loc(#loc4326)
      %36246 = stablehlo.broadcast_in_dim %36245, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4327)
      %36247 = func.call @_where_906(%30549#1, %36246) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36248 = stablehlo.multiply %33846, %36239#0 : tensor<1x1x2048x2048xbf16> loc(#loc3392)
      %36249 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4328)
      %36250 = stablehlo.reduce(%36248 init: %36249) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4328)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4328)
      } loc(#loc4328)
      %36251 = stablehlo.reshape %36250 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4329)
      %36252 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4330)
      %36253 = stablehlo.reduce(%36251 init: %36252) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4330)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4330)
      } loc(#loc4330)
      %36254 = stablehlo.broadcast_in_dim %36253, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4331)
      %36255 = func.call @_where_907(%30550#1, %36254) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36256 = stablehlo.broadcast_in_dim %33850, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3392)
      %36257 = stablehlo.multiply %36239#0, %36256 : tensor<1x1x2048x2048xbf16> loc(#loc3392)
      %36258 = stablehlo.multiply %33836, %36257 : tensor<1x1x2048x2048xbf16> loc(#loc3389)
      %36259 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4332)
      %36260 = stablehlo.reduce(%36258 init: %36259) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4332)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4332)
      } loc(#loc4332)
      %36261 = stablehlo.reshape %36260 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4333)
      %36262 = stablehlo.multiply %36261, %33844 : tensor<1x1x2048x1xbf16> loc(#loc3389)
      %36263 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36264 = stablehlo.broadcast_in_dim %36263, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3383)
      %36265 = stablehlo.divide %36262, %36264 : tensor<1x1x2048x1xbf16> loc(#loc3383)
      %36266 = stablehlo.convert %36265 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3379)
      %36267 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4334)
      %36268 = stablehlo.reduce(%36266 init: %36267) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4334)
        stablehlo.return %39096 : tensor<f32> loc(#loc4334)
      } loc(#loc4334)
      %36269 = stablehlo.broadcast_in_dim %36268, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4335)
      %36270 = stablehlo.broadcast_in_dim %36269, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4336)
      %36271 = stablehlo.convert %36270 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3382)
      %36272 = stablehlo.multiply %36271, %33826 : tensor<1x1x2048x2048xbf16> loc(#loc3385)
      %36273 = stablehlo.negate %36272 : tensor<1x1x2048x2048xbf16> loc(#loc4337)
      %36274 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4338)
      %36275 = stablehlo.reduce(%36273 init: %36274) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4338)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4338)
      } loc(#loc4338)
      %36276 = stablehlo.reshape %36275 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4339)
      %36277 = stablehlo.add %36177, %36272 : tensor<1x1x2048x2048xbf16> loc(#loc4340)
      %36278 = stablehlo.broadcast_in_dim %33840, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3389)
      %36279 = stablehlo.multiply %36257, %36278 : tensor<1x1x2048x2048xbf16> loc(#loc3389)
      %36280 = stablehlo.negate %36279 : tensor<1x1x2048x2048xbf16> loc(#loc4337)
      %36281 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4338)
      %36282 = stablehlo.reduce(%36280 init: %36281) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4338)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4338)
      } loc(#loc4338)
      %36283 = stablehlo.reshape %36282 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4339)
      %36284 = stablehlo.add %36277, %36279 : tensor<1x1x2048x2048xbf16> loc(#loc4340)
      %36285 = stablehlo.add %36276, %36283 : tensor<1x1x2048x1xbf16> loc(#loc4340)
      %36286 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36287 = stablehlo.broadcast_in_dim %36286, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3383)
      %36288 = stablehlo.divide %36285, %36287 : tensor<1x1x2048x1xbf16> loc(#loc3383)
      %36289 = stablehlo.convert %36288 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3379)
      %36290 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4334)
      %36291 = stablehlo.reduce(%36289 init: %36290) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4334)
        stablehlo.return %39096 : tensor<f32> loc(#loc4334)
      } loc(#loc4334)
      %36292 = stablehlo.broadcast_in_dim %36291, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4335)
      %36293 = stablehlo.broadcast_in_dim %36292, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4336)
      %36294 = stablehlo.convert %36293 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3382)
      %36295 = stablehlo.add %36284, %36294 : tensor<1x1x2048x2048xbf16> loc(#loc4341)
      %36296 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %36297 = stablehlo.broadcast_in_dim %36296, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3354)
      %36298 = stablehlo.multiply %36295, %36297 : tensor<1x1x2048x2048xbf16> loc(#loc3354)
      %36299 = stablehlo.broadcast_in_dim %33748, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3353)
      %36300 = stablehlo.multiply %36298, %36299 : tensor<1x1x2048x2048xbf16> loc(#loc3353)
      %36301 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4342)
      %36302 = stablehlo.reduce(%36300 init: %36301) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4342)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4342)
      } loc(#loc4342)
      %36303 = stablehlo.reshape %36302 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4343)
      %36304 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4344)
      %36305 = stablehlo.reduce(%36303 init: %36304) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4344)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4344)
      } loc(#loc4344)
      %36306 = stablehlo.broadcast_in_dim %36305, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4345)
      %36307 = func.call @_where_908(%30534#1, %36306) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36308 = arith.select %ptrue, %36300, %36300 : tensor<1x1x2048x2048xbf16>
      %36309:2 = func.call @_einsum_909(%30535#0, %33738, %36308) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4346)
      %36310 = func.call @_where_910(%30535#1, %36309#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %36311 = stablehlo.broadcast_in_dim %33736, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3347)
      %36312 = stablehlo.multiply %36309#0, %36311 : tensor<1x1x2048x8192xbf16> loc(#loc3347)
      %36313 = stablehlo.multiply %33707, %36312 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36314 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %36315 = stablehlo.broadcast_in_dim %36314, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36316 = stablehlo.multiply %36315, %36313 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36317 = stablehlo.multiply %36316, %33724 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36318 = stablehlo.multiply %36317, %33721 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36319 = stablehlo.add %36317, %36318 : tensor<1x1x2048x8192xbf16> loc(#loc4347)
      %36320 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %36321 = stablehlo.broadcast_in_dim %36320, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36322 = stablehlo.multiply %36321, %36319 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36323 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %36324 = stablehlo.broadcast_in_dim %36323, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36325 = stablehlo.multiply %36324, %36322 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36326 = stablehlo.multiply %36325, %33713 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36327 = stablehlo.add %36322, %36326 : tensor<1x1x2048x8192xbf16> loc(#loc4347)
      %36328 = stablehlo.multiply %36312, %33730 : tensor<1x1x2048x8192xbf16> loc(#loc3340)
      %36329 = stablehlo.add %36327, %36328 : tensor<1x1x2048x8192xbf16> loc(#loc4347)
      %36330 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4348)
      %36331 = stablehlo.reduce(%36329 init: %36330) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4348)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4348)
      } loc(#loc4348)
      %36332 = stablehlo.reshape %36331 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4349)
      %36333 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4350)
      %36334 = stablehlo.reduce(%36332 init: %36333) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4350)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4350)
      } loc(#loc4350)
      %36335 = stablehlo.broadcast_in_dim %36334, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4351)
      %36336 = func.call @_where_911(%30532#1, %36335) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %36337 = arith.select %ptrue, %36329, %36329 : tensor<1x1x2048x8192xbf16>
      %36338:2 = func.call @_einsum_912(%30533#0, %33702, %36337) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4352)
      %36339 = func.call @_where_913(%30533#1, %36338#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %36340 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4353)
      %36341 = stablehlo.reduce(%36338#0 init: %36340) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4353)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4353)
      } loc(#loc4353)
      %36342 = stablehlo.reshape %36341 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4354)
      %36343 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4355)
      %36344 = stablehlo.reduce(%36342 init: %36343) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4355)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4355)
      } loc(#loc4355)
      %36345 = stablehlo.broadcast_in_dim %36344, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4356)
      %36346 = func.call @_where_914(%30536#1, %36345) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36347 = stablehlo.multiply %33693, %36338#0 : tensor<1x1x2048x2048xbf16> loc(#loc3333)
      %36348 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4357)
      %36349 = stablehlo.reduce(%36347 init: %36348) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4357)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4357)
      } loc(#loc4357)
      %36350 = stablehlo.reshape %36349 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4358)
      %36351 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4359)
      %36352 = stablehlo.reduce(%36350 init: %36351) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4359)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4359)
      } loc(#loc4359)
      %36353 = stablehlo.broadcast_in_dim %36352, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4360)
      %36354 = func.call @_where_915(%30537#1, %36353) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36355 = stablehlo.broadcast_in_dim %33697, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3333)
      %36356 = stablehlo.multiply %36338#0, %36355 : tensor<1x1x2048x2048xbf16> loc(#loc3333)
      %36357 = stablehlo.multiply %33683, %36356 : tensor<1x1x2048x2048xbf16> loc(#loc3330)
      %36358 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4361)
      %36359 = stablehlo.reduce(%36357 init: %36358) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4361)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4361)
      } loc(#loc4361)
      %36360 = stablehlo.reshape %36359 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4362)
      %36361 = stablehlo.multiply %36360, %33691 : tensor<1x1x2048x1xbf16> loc(#loc3330)
      %36362 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36363 = stablehlo.broadcast_in_dim %36362, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3324)
      %36364 = stablehlo.divide %36361, %36363 : tensor<1x1x2048x1xbf16> loc(#loc3324)
      %36365 = stablehlo.convert %36364 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3320)
      %36366 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4363)
      %36367 = stablehlo.reduce(%36365 init: %36366) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4363)
        stablehlo.return %39096 : tensor<f32> loc(#loc4363)
      } loc(#loc4363)
      %36368 = stablehlo.broadcast_in_dim %36367, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4364)
      %36369 = stablehlo.broadcast_in_dim %36368, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4365)
      %36370 = stablehlo.convert %36369 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3323)
      %36371 = stablehlo.multiply %36370, %33673 : tensor<1x1x2048x2048xbf16> loc(#loc3326)
      %36372 = stablehlo.negate %36371 : tensor<1x1x2048x2048xbf16> loc(#loc4366)
      %36373 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4367)
      %36374 = stablehlo.reduce(%36372 init: %36373) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4367)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4367)
      } loc(#loc4367)
      %36375 = stablehlo.reshape %36374 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4368)
      %36376 = stablehlo.add %36295, %36371 : tensor<1x1x2048x2048xbf16> loc(#loc4369)
      %36377 = stablehlo.broadcast_in_dim %33687, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3330)
      %36378 = stablehlo.multiply %36356, %36377 : tensor<1x1x2048x2048xbf16> loc(#loc3330)
      %36379 = stablehlo.negate %36378 : tensor<1x1x2048x2048xbf16> loc(#loc4366)
      %36380 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4367)
      %36381 = stablehlo.reduce(%36379 init: %36380) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4367)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4367)
      } loc(#loc4367)
      %36382 = stablehlo.reshape %36381 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4368)
      %36383 = stablehlo.add %36376, %36378 : tensor<1x1x2048x2048xbf16> loc(#loc4369)
      %36384 = stablehlo.add %36375, %36382 : tensor<1x1x2048x1xbf16> loc(#loc4369)
      %36385 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36386 = stablehlo.broadcast_in_dim %36385, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3324)
      %36387 = stablehlo.divide %36384, %36386 : tensor<1x1x2048x1xbf16> loc(#loc3324)
      %36388 = stablehlo.convert %36387 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3320)
      %36389 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4363)
      %36390 = stablehlo.reduce(%36388 init: %36389) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4363)
        stablehlo.return %39096 : tensor<f32> loc(#loc4363)
      } loc(#loc4363)
      %36391 = stablehlo.broadcast_in_dim %36390, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4364)
      %36392 = stablehlo.broadcast_in_dim %36391, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4365)
      %36393 = stablehlo.convert %36392 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3323)
      %36394 = stablehlo.add %36383, %36393 : tensor<1x1x2048x2048xbf16> loc(#loc4370)
      %36395 = stablehlo.add %35296, %36394 : tensor<1x1x2048x2048xbf16> loc(#loc4371)
      %36396 = arith.select %ptrue, %36394, %36394 : tensor<1x1x2048x2048xbf16>
      %36397:2 = func.call @_einsum_916(%30542#0, %33655, %36396) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4372)
      %36398 = func.call @_where_917(%30542#1, %36397#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %36399 = arith.select %ptrue, %36397#0, %36397#0 : tensor<1x1x2048x32x64xbf16>
      %36400:2 = func.call @_einsum_918(%33653, %33610, %36399) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4373)
      %36401 = arith.select %ptrue, %36400#1, %36400#1 : tensor<1x1x2048x32x64xbf16>
      %36402 = stablehlo.broadcast_in_dim %36401, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4374)
      %36403 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36404 = stablehlo.pad %36402, %36403, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4375)
      %36405 = stablehlo.convert %36400#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4376)
      %36406 = stablehlo.broadcast_in_dim %33652, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3311)
      %36407 = stablehlo.multiply %36405, %36406 : tensor<1x1x32x2048x2048xf32> loc(#loc3311)
      %36408 = stablehlo.multiply %36407, %33643 : tensor<1x1x32x2048x2048xf32> loc(#loc3311)
      %36409 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3309)
      %36410 = stablehlo.reduce(%36408 init: %36409) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3309)
        stablehlo.return %39096 : tensor<f32> loc(#loc3309)
      } loc(#loc3309)
      %36411 = stablehlo.reshape %36410 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4377)
      %36412 = stablehlo.negate %36411 : tensor<1x1x32x2048x1xf32> loc(#loc4378)
      %36413 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4379)
      %36414 = stablehlo.reduce(%36412 init: %36413) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4379)
        stablehlo.return %39096 : tensor<f32> loc(#loc4379)
      } loc(#loc4379)
      %36415 = stablehlo.broadcast_in_dim %36414, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4380)
      %36416 = stablehlo.broadcast_in_dim %36415, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4381)
      %36417 = stablehlo.broadcast_in_dim %33646, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3310)
      %36418 = stablehlo.divide %36405, %36417 : tensor<1x1x32x2048x2048xf32> loc(#loc3310)
      %36419 = stablehlo.add %36416, %36418 : tensor<1x1x32x2048x2048xf32> loc(#loc4382)
      %36420 = stablehlo.multiply %36419, %33643 : tensor<1x1x32x2048x2048xf32> loc(#loc3311)
      %36421 = func.call @_where_919(%33637#1, %36420) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4383)
      %36422 = stablehlo.convert %36421 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4384)
      %36423 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36424 = stablehlo.broadcast_in_dim %36423, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3299)
      %36425 = stablehlo.multiply %36424, %36422 : tensor<1x1x32x2048x2048xbf16> loc(#loc3299)
      %36426 = stablehlo.multiply %36425, %33625 : tensor<1x1x32x2048x2048xbf16> loc(#loc3299)
      %36427 = stablehlo.multiply %36426, %33622 : tensor<1x1x32x2048x2048xbf16> loc(#loc3299)
      %36428 = stablehlo.add %36426, %36427 : tensor<1x1x32x2048x2048xbf16> loc(#loc4385)
      %36429 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36430 = stablehlo.broadcast_in_dim %36429, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3296)
      %36431 = stablehlo.divide %36428, %36430 : tensor<1x1x32x2048x2048xbf16> loc(#loc3296)
      %36432:2 = func.call @_einsum_920(%33609, %33617, %36431) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4386)
      %36433 = arith.select %ptrue, %36432#1, %36432#1 : tensor<1x1x2048x32x64xbf16>
      %36434 = stablehlo.broadcast_in_dim %36433, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4374)
      %36435 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36436 = stablehlo.pad %36434, %36435, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4387)
      %36437 = stablehlo.add %36404, %36436 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4388)
      %36438 = stablehlo.multiply %33608, %36432#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3294)
      %36439 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4389)
      %36440 = stablehlo.reduce(%36438 init: %36439) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4389)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4389)
      } loc(#loc4389)
      %36441 = stablehlo.reshape %36440 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4390)
      %36442 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4391)
      %36443 = stablehlo.reduce(%36441 init: %36442) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4391)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4391)
      } loc(#loc4391)
      %36444 = stablehlo.broadcast_in_dim %36443, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4392)
      %36445 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %36446 = stablehlo.broadcast_in_dim %36445, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3292)
      %36447 = stablehlo.multiply %36446, %36444 : tensor<1x64xbf16> loc(#loc3292)
      %36448 = func.call @softplus_921(%33611#1, %36447) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4393)
      %36449 = func.call @_where_922(%30541#1, %36448) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %36450 = stablehlo.broadcast_in_dim %33615, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3294)
      %36451 = stablehlo.multiply %36432#0, %36450 : tensor<1x1x2048x32x64xbf16> loc(#loc3294)
      %36452 = arith.select %ptrue, %36451, %36451 : tensor<1x1x2048x32x64xbf16>
      %36453 = stablehlo.broadcast_in_dim %36452, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4374)
      %36454 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36455 = stablehlo.pad %36453, %36454, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4394)
      %36456 = stablehlo.add %36437, %36455 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4388)
      %36457:2 = func.call @_einsum_923(%33600, %30540#0, %36456) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4395)
      %36458 = func.call @_where_924(%30540#1, %36457#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %36459 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4396)
      %36460 = stablehlo.reduce(%36457#0 init: %36459) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4396)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4396)
      } loc(#loc4396)
      %36461 = stablehlo.reshape %36460 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4397)
      %36462 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4398)
      %36463 = stablehlo.reduce(%36461 init: %36462) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4398)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4398)
      } loc(#loc4398)
      %36464 = stablehlo.broadcast_in_dim %36463, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4399)
      %36465 = func.call @_where_925(%30538#1, %36464) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36466 = stablehlo.multiply %33591, %36457#0 : tensor<1x1x2048x2048xbf16> loc(#loc3282)
      %36467 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4400)
      %36468 = stablehlo.reduce(%36466 init: %36467) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4400)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4400)
      } loc(#loc4400)
      %36469 = stablehlo.reshape %36468 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4401)
      %36470 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4402)
      %36471 = stablehlo.reduce(%36469 init: %36470) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4402)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4402)
      } loc(#loc4402)
      %36472 = stablehlo.broadcast_in_dim %36471, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4403)
      %36473 = func.call @_where_926(%30539#1, %36472) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36474 = stablehlo.broadcast_in_dim %33595, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3282)
      %36475 = stablehlo.multiply %36457#0, %36474 : tensor<1x1x2048x2048xbf16> loc(#loc3282)
      %36476 = stablehlo.multiply %33581, %36475 : tensor<1x1x2048x2048xbf16> loc(#loc3279)
      %36477 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4404)
      %36478 = stablehlo.reduce(%36476 init: %36477) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4404)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4404)
      } loc(#loc4404)
      %36479 = stablehlo.reshape %36478 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4405)
      %36480 = stablehlo.multiply %36479, %33589 : tensor<1x1x2048x1xbf16> loc(#loc3279)
      %36481 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36482 = stablehlo.broadcast_in_dim %36481, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3273)
      %36483 = stablehlo.divide %36480, %36482 : tensor<1x1x2048x1xbf16> loc(#loc3273)
      %36484 = stablehlo.convert %36483 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3269)
      %36485 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4406)
      %36486 = stablehlo.reduce(%36484 init: %36485) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4406)
        stablehlo.return %39096 : tensor<f32> loc(#loc4406)
      } loc(#loc4406)
      %36487 = stablehlo.broadcast_in_dim %36486, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4407)
      %36488 = stablehlo.broadcast_in_dim %36487, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4408)
      %36489 = stablehlo.convert %36488 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3272)
      %36490 = stablehlo.multiply %36489, %33571 : tensor<1x1x2048x2048xbf16> loc(#loc3275)
      %36491 = stablehlo.negate %36490 : tensor<1x1x2048x2048xbf16> loc(#loc4409)
      %36492 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4410)
      %36493 = stablehlo.reduce(%36491 init: %36492) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4410)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4410)
      } loc(#loc4410)
      %36494 = stablehlo.reshape %36493 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4411)
      %36495 = stablehlo.add %36395, %36490 : tensor<1x1x2048x2048xbf16> loc(#loc4412)
      %36496 = stablehlo.broadcast_in_dim %33585, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3279)
      %36497 = stablehlo.multiply %36475, %36496 : tensor<1x1x2048x2048xbf16> loc(#loc3279)
      %36498 = stablehlo.negate %36497 : tensor<1x1x2048x2048xbf16> loc(#loc4409)
      %36499 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4410)
      %36500 = stablehlo.reduce(%36498 init: %36499) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4410)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4410)
      } loc(#loc4410)
      %36501 = stablehlo.reshape %36500 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4411)
      %36502 = stablehlo.add %36495, %36497 : tensor<1x1x2048x2048xbf16> loc(#loc4412)
      %36503 = stablehlo.add %36494, %36501 : tensor<1x1x2048x1xbf16> loc(#loc4412)
      %36504 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36505 = stablehlo.broadcast_in_dim %36504, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3273)
      %36506 = stablehlo.divide %36503, %36505 : tensor<1x1x2048x1xbf16> loc(#loc3273)
      %36507 = stablehlo.convert %36506 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3269)
      %36508 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4406)
      %36509 = stablehlo.reduce(%36507 init: %36508) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4406)
        stablehlo.return %39096 : tensor<f32> loc(#loc4406)
      } loc(#loc4406)
      %36510 = stablehlo.broadcast_in_dim %36509, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4407)
      %36511 = stablehlo.broadcast_in_dim %36510, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4408)
      %36512 = stablehlo.convert %36511 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3272)
      %36513 = stablehlo.add %36502, %36512 : tensor<1x1x2048x2048xbf16> loc(#loc4413)
      %36514 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %36515 = stablehlo.broadcast_in_dim %36514, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3244)
      %36516 = stablehlo.multiply %36513, %36515 : tensor<1x1x2048x2048xbf16> loc(#loc3244)
      %36517 = stablehlo.broadcast_in_dim %33493, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3243)
      %36518 = stablehlo.multiply %36516, %36517 : tensor<1x1x2048x2048xbf16> loc(#loc3243)
      %36519 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4414)
      %36520 = stablehlo.reduce(%36518 init: %36519) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4414)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4414)
      } loc(#loc4414)
      %36521 = stablehlo.reshape %36520 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4415)
      %36522 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4416)
      %36523 = stablehlo.reduce(%36521 init: %36522) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4416)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4416)
      } loc(#loc4416)
      %36524 = stablehlo.broadcast_in_dim %36523, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4417)
      %36525 = func.call @_where_927(%30523#1, %36524) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36526 = arith.select %ptrue, %36518, %36518 : tensor<1x1x2048x2048xbf16>
      %36527:2 = func.call @_einsum_928(%30524#0, %33483, %36526) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4418)
      %36528 = func.call @_where_929(%30524#1, %36527#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %36529 = stablehlo.broadcast_in_dim %33481, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3237)
      %36530 = stablehlo.multiply %36527#0, %36529 : tensor<1x1x2048x8192xbf16> loc(#loc3237)
      %36531 = stablehlo.multiply %33452, %36530 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36532 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %36533 = stablehlo.broadcast_in_dim %36532, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36534 = stablehlo.multiply %36533, %36531 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36535 = stablehlo.multiply %36534, %33469 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36536 = stablehlo.multiply %36535, %33466 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36537 = stablehlo.add %36535, %36536 : tensor<1x1x2048x8192xbf16> loc(#loc4419)
      %36538 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %36539 = stablehlo.broadcast_in_dim %36538, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36540 = stablehlo.multiply %36539, %36537 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36541 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %36542 = stablehlo.broadcast_in_dim %36541, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36543 = stablehlo.multiply %36542, %36540 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36544 = stablehlo.multiply %36543, %33458 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36545 = stablehlo.add %36540, %36544 : tensor<1x1x2048x8192xbf16> loc(#loc4419)
      %36546 = stablehlo.multiply %36530, %33475 : tensor<1x1x2048x8192xbf16> loc(#loc3230)
      %36547 = stablehlo.add %36545, %36546 : tensor<1x1x2048x8192xbf16> loc(#loc4419)
      %36548 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4420)
      %36549 = stablehlo.reduce(%36547 init: %36548) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4420)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4420)
      } loc(#loc4420)
      %36550 = stablehlo.reshape %36549 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4421)
      %36551 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4422)
      %36552 = stablehlo.reduce(%36550 init: %36551) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4422)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4422)
      } loc(#loc4422)
      %36553 = stablehlo.broadcast_in_dim %36552, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4423)
      %36554 = func.call @_where_930(%30521#1, %36553) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %36555 = arith.select %ptrue, %36547, %36547 : tensor<1x1x2048x8192xbf16>
      %36556:2 = func.call @_einsum_931(%30522#0, %33447, %36555) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4424)
      %36557 = func.call @_where_932(%30522#1, %36556#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %36558 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4425)
      %36559 = stablehlo.reduce(%36556#0 init: %36558) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4425)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4425)
      } loc(#loc4425)
      %36560 = stablehlo.reshape %36559 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4426)
      %36561 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4427)
      %36562 = stablehlo.reduce(%36560 init: %36561) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4427)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4427)
      } loc(#loc4427)
      %36563 = stablehlo.broadcast_in_dim %36562, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4428)
      %36564 = func.call @_where_933(%30525#1, %36563) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36565 = stablehlo.multiply %33438, %36556#0 : tensor<1x1x2048x2048xbf16> loc(#loc3223)
      %36566 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4429)
      %36567 = stablehlo.reduce(%36565 init: %36566) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4429)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4429)
      } loc(#loc4429)
      %36568 = stablehlo.reshape %36567 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4430)
      %36569 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4431)
      %36570 = stablehlo.reduce(%36568 init: %36569) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4431)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4431)
      } loc(#loc4431)
      %36571 = stablehlo.broadcast_in_dim %36570, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4432)
      %36572 = func.call @_where_934(%30526#1, %36571) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36573 = stablehlo.broadcast_in_dim %33442, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3223)
      %36574 = stablehlo.multiply %36556#0, %36573 : tensor<1x1x2048x2048xbf16> loc(#loc3223)
      %36575 = stablehlo.multiply %33428, %36574 : tensor<1x1x2048x2048xbf16> loc(#loc3220)
      %36576 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4433)
      %36577 = stablehlo.reduce(%36575 init: %36576) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4433)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4433)
      } loc(#loc4433)
      %36578 = stablehlo.reshape %36577 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4434)
      %36579 = stablehlo.multiply %36578, %33436 : tensor<1x1x2048x1xbf16> loc(#loc3220)
      %36580 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36581 = stablehlo.broadcast_in_dim %36580, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3214)
      %36582 = stablehlo.divide %36579, %36581 : tensor<1x1x2048x1xbf16> loc(#loc3214)
      %36583 = stablehlo.convert %36582 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3210)
      %36584 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4435)
      %36585 = stablehlo.reduce(%36583 init: %36584) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4435)
        stablehlo.return %39096 : tensor<f32> loc(#loc4435)
      } loc(#loc4435)
      %36586 = stablehlo.broadcast_in_dim %36585, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4436)
      %36587 = stablehlo.broadcast_in_dim %36586, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4437)
      %36588 = stablehlo.convert %36587 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3213)
      %36589 = stablehlo.multiply %36588, %33418 : tensor<1x1x2048x2048xbf16> loc(#loc3216)
      %36590 = stablehlo.negate %36589 : tensor<1x1x2048x2048xbf16> loc(#loc4438)
      %36591 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4439)
      %36592 = stablehlo.reduce(%36590 init: %36591) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4439)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4439)
      } loc(#loc4439)
      %36593 = stablehlo.reshape %36592 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4440)
      %36594 = stablehlo.add %36513, %36589 : tensor<1x1x2048x2048xbf16> loc(#loc4441)
      %36595 = stablehlo.broadcast_in_dim %33432, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3220)
      %36596 = stablehlo.multiply %36574, %36595 : tensor<1x1x2048x2048xbf16> loc(#loc3220)
      %36597 = stablehlo.negate %36596 : tensor<1x1x2048x2048xbf16> loc(#loc4438)
      %36598 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4439)
      %36599 = stablehlo.reduce(%36597 init: %36598) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4439)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4439)
      } loc(#loc4439)
      %36600 = stablehlo.reshape %36599 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4440)
      %36601 = stablehlo.add %36594, %36596 : tensor<1x1x2048x2048xbf16> loc(#loc4441)
      %36602 = stablehlo.add %36593, %36600 : tensor<1x1x2048x1xbf16> loc(#loc4441)
      %36603 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36604 = stablehlo.broadcast_in_dim %36603, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3214)
      %36605 = stablehlo.divide %36602, %36604 : tensor<1x1x2048x1xbf16> loc(#loc3214)
      %36606 = stablehlo.convert %36605 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3210)
      %36607 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4435)
      %36608 = stablehlo.reduce(%36606 init: %36607) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4435)
        stablehlo.return %39096 : tensor<f32> loc(#loc4435)
      } loc(#loc4435)
      %36609 = stablehlo.broadcast_in_dim %36608, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4436)
      %36610 = stablehlo.broadcast_in_dim %36609, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4437)
      %36611 = stablehlo.convert %36610 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3213)
      %36612 = stablehlo.add %36601, %36611 : tensor<1x1x2048x2048xbf16> loc(#loc4442)
      %36613 = stablehlo.add %35337, %36612 : tensor<1x1x2048x2048xbf16> loc(#loc4443)
      %36614 = arith.select %ptrue, %36612, %36612 : tensor<1x1x2048x2048xbf16>
      %36615:2 = func.call @_einsum_935(%30531#0, %33400, %36614) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4444)
      %36616 = func.call @_where_936(%30531#1, %36615#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %36617 = arith.select %ptrue, %36615#0, %36615#0 : tensor<1x1x2048x32x64xbf16>
      %36618:2 = func.call @_einsum_937(%33398, %33355, %36617) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4445)
      %36619 = arith.select %ptrue, %36618#1, %36618#1 : tensor<1x1x2048x32x64xbf16>
      %36620 = stablehlo.broadcast_in_dim %36619, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4446)
      %36621 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36622 = stablehlo.pad %36620, %36621, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4447)
      %36623 = stablehlo.convert %36618#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4448)
      %36624 = stablehlo.broadcast_in_dim %33397, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3201)
      %36625 = stablehlo.multiply %36623, %36624 : tensor<1x1x32x2048x2048xf32> loc(#loc3201)
      %36626 = stablehlo.multiply %36625, %33388 : tensor<1x1x32x2048x2048xf32> loc(#loc3201)
      %36627 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3199)
      %36628 = stablehlo.reduce(%36626 init: %36627) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3199)
        stablehlo.return %39096 : tensor<f32> loc(#loc3199)
      } loc(#loc3199)
      %36629 = stablehlo.reshape %36628 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4449)
      %36630 = stablehlo.negate %36629 : tensor<1x1x32x2048x1xf32> loc(#loc4450)
      %36631 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4451)
      %36632 = stablehlo.reduce(%36630 init: %36631) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4451)
        stablehlo.return %39096 : tensor<f32> loc(#loc4451)
      } loc(#loc4451)
      %36633 = stablehlo.broadcast_in_dim %36632, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4452)
      %36634 = stablehlo.broadcast_in_dim %36633, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4453)
      %36635 = stablehlo.broadcast_in_dim %33391, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3200)
      %36636 = stablehlo.divide %36623, %36635 : tensor<1x1x32x2048x2048xf32> loc(#loc3200)
      %36637 = stablehlo.add %36634, %36636 : tensor<1x1x32x2048x2048xf32> loc(#loc4454)
      %36638 = stablehlo.multiply %36637, %33388 : tensor<1x1x32x2048x2048xf32> loc(#loc3201)
      %36639 = func.call @_where_938(%33382#1, %36638) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4455)
      %36640 = stablehlo.convert %36639 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4456)
      %36641 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36642 = stablehlo.broadcast_in_dim %36641, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3189)
      %36643 = stablehlo.multiply %36642, %36640 : tensor<1x1x32x2048x2048xbf16> loc(#loc3189)
      %36644 = stablehlo.multiply %36643, %33370 : tensor<1x1x32x2048x2048xbf16> loc(#loc3189)
      %36645 = stablehlo.multiply %36644, %33367 : tensor<1x1x32x2048x2048xbf16> loc(#loc3189)
      %36646 = stablehlo.add %36644, %36645 : tensor<1x1x32x2048x2048xbf16> loc(#loc4457)
      %36647 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36648 = stablehlo.broadcast_in_dim %36647, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3186)
      %36649 = stablehlo.divide %36646, %36648 : tensor<1x1x32x2048x2048xbf16> loc(#loc3186)
      %36650:2 = func.call @_einsum_939(%33354, %33362, %36649) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4458)
      %36651 = arith.select %ptrue, %36650#1, %36650#1 : tensor<1x1x2048x32x64xbf16>
      %36652 = stablehlo.broadcast_in_dim %36651, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4446)
      %36653 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36654 = stablehlo.pad %36652, %36653, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4459)
      %36655 = stablehlo.add %36622, %36654 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4460)
      %36656 = stablehlo.multiply %33353, %36650#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3184)
      %36657 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4461)
      %36658 = stablehlo.reduce(%36656 init: %36657) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4461)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4461)
      } loc(#loc4461)
      %36659 = stablehlo.reshape %36658 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4462)
      %36660 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4463)
      %36661 = stablehlo.reduce(%36659 init: %36660) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4463)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4463)
      } loc(#loc4463)
      %36662 = stablehlo.broadcast_in_dim %36661, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4464)
      %36663 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %36664 = stablehlo.broadcast_in_dim %36663, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3182)
      %36665 = stablehlo.multiply %36664, %36662 : tensor<1x64xbf16> loc(#loc3182)
      %36666 = func.call @softplus_940(%33356#1, %36665) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4465)
      %36667 = func.call @_where_941(%30530#1, %36666) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %36668 = stablehlo.broadcast_in_dim %33360, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3184)
      %36669 = stablehlo.multiply %36650#0, %36668 : tensor<1x1x2048x32x64xbf16> loc(#loc3184)
      %36670 = arith.select %ptrue, %36669, %36669 : tensor<1x1x2048x32x64xbf16>
      %36671 = stablehlo.broadcast_in_dim %36670, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4446)
      %36672 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36673 = stablehlo.pad %36671, %36672, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4466)
      %36674 = stablehlo.add %36655, %36673 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4460)
      %36675:2 = func.call @_einsum_942(%33345, %30529#0, %36674) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4467)
      %36676 = func.call @_where_943(%30529#1, %36675#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %36677 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4468)
      %36678 = stablehlo.reduce(%36675#0 init: %36677) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4468)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4468)
      } loc(#loc4468)
      %36679 = stablehlo.reshape %36678 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4469)
      %36680 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4470)
      %36681 = stablehlo.reduce(%36679 init: %36680) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4470)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4470)
      } loc(#loc4470)
      %36682 = stablehlo.broadcast_in_dim %36681, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4471)
      %36683 = func.call @_where_944(%30527#1, %36682) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36684 = stablehlo.multiply %33336, %36675#0 : tensor<1x1x2048x2048xbf16> loc(#loc3172)
      %36685 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4472)
      %36686 = stablehlo.reduce(%36684 init: %36685) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4472)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4472)
      } loc(#loc4472)
      %36687 = stablehlo.reshape %36686 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4473)
      %36688 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4474)
      %36689 = stablehlo.reduce(%36687 init: %36688) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4474)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4474)
      } loc(#loc4474)
      %36690 = stablehlo.broadcast_in_dim %36689, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4475)
      %36691 = func.call @_where_945(%30528#1, %36690) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36692 = stablehlo.broadcast_in_dim %33340, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3172)
      %36693 = stablehlo.multiply %36675#0, %36692 : tensor<1x1x2048x2048xbf16> loc(#loc3172)
      %36694 = stablehlo.multiply %33326, %36693 : tensor<1x1x2048x2048xbf16> loc(#loc3169)
      %36695 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4476)
      %36696 = stablehlo.reduce(%36694 init: %36695) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4476)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4476)
      } loc(#loc4476)
      %36697 = stablehlo.reshape %36696 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4477)
      %36698 = stablehlo.multiply %36697, %33334 : tensor<1x1x2048x1xbf16> loc(#loc3169)
      %36699 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36700 = stablehlo.broadcast_in_dim %36699, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3163)
      %36701 = stablehlo.divide %36698, %36700 : tensor<1x1x2048x1xbf16> loc(#loc3163)
      %36702 = stablehlo.convert %36701 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3159)
      %36703 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4478)
      %36704 = stablehlo.reduce(%36702 init: %36703) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4478)
        stablehlo.return %39096 : tensor<f32> loc(#loc4478)
      } loc(#loc4478)
      %36705 = stablehlo.broadcast_in_dim %36704, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4479)
      %36706 = stablehlo.broadcast_in_dim %36705, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4480)
      %36707 = stablehlo.convert %36706 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3162)
      %36708 = stablehlo.multiply %36707, %33316 : tensor<1x1x2048x2048xbf16> loc(#loc3165)
      %36709 = stablehlo.negate %36708 : tensor<1x1x2048x2048xbf16> loc(#loc4481)
      %36710 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4482)
      %36711 = stablehlo.reduce(%36709 init: %36710) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4482)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4482)
      } loc(#loc4482)
      %36712 = stablehlo.reshape %36711 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4483)
      %36713 = stablehlo.add %36613, %36708 : tensor<1x1x2048x2048xbf16> loc(#loc4484)
      %36714 = stablehlo.broadcast_in_dim %33330, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3169)
      %36715 = stablehlo.multiply %36693, %36714 : tensor<1x1x2048x2048xbf16> loc(#loc3169)
      %36716 = stablehlo.negate %36715 : tensor<1x1x2048x2048xbf16> loc(#loc4481)
      %36717 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4482)
      %36718 = stablehlo.reduce(%36716 init: %36717) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4482)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4482)
      } loc(#loc4482)
      %36719 = stablehlo.reshape %36718 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4483)
      %36720 = stablehlo.add %36713, %36715 : tensor<1x1x2048x2048xbf16> loc(#loc4484)
      %36721 = stablehlo.add %36712, %36719 : tensor<1x1x2048x1xbf16> loc(#loc4484)
      %36722 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36723 = stablehlo.broadcast_in_dim %36722, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3163)
      %36724 = stablehlo.divide %36721, %36723 : tensor<1x1x2048x1xbf16> loc(#loc3163)
      %36725 = stablehlo.convert %36724 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3159)
      %36726 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4478)
      %36727 = stablehlo.reduce(%36725 init: %36726) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4478)
        stablehlo.return %39096 : tensor<f32> loc(#loc4478)
      } loc(#loc4478)
      %36728 = stablehlo.broadcast_in_dim %36727, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4479)
      %36729 = stablehlo.broadcast_in_dim %36728, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4480)
      %36730 = stablehlo.convert %36729 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3162)
      %36731 = stablehlo.add %36720, %36730 : tensor<1x1x2048x2048xbf16> loc(#loc4485)
      %36732 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %36733 = stablehlo.broadcast_in_dim %36732, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3134)
      %36734 = stablehlo.multiply %36731, %36733 : tensor<1x1x2048x2048xbf16> loc(#loc3134)
      %36735 = stablehlo.broadcast_in_dim %33238, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3133)
      %36736 = stablehlo.multiply %36734, %36735 : tensor<1x1x2048x2048xbf16> loc(#loc3133)
      %36737 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4486)
      %36738 = stablehlo.reduce(%36736 init: %36737) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4486)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4486)
      } loc(#loc4486)
      %36739 = stablehlo.reshape %36738 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4487)
      %36740 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4488)
      %36741 = stablehlo.reduce(%36739 init: %36740) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4488)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4488)
      } loc(#loc4488)
      %36742 = stablehlo.broadcast_in_dim %36741, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4489)
      %36743 = func.call @_where_946(%30666#1, %36742) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36744 = arith.select %ptrue, %36736, %36736 : tensor<1x1x2048x2048xbf16>
      %36745:2 = func.call @_einsum_947(%30667#0, %33228, %36744) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4490)
      %36746 = func.call @_where_948(%30667#1, %36745#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %36747 = stablehlo.broadcast_in_dim %33226, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3127)
      %36748 = stablehlo.multiply %36745#0, %36747 : tensor<1x1x2048x8192xbf16> loc(#loc3127)
      %36749 = stablehlo.multiply %33197, %36748 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36750 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %36751 = stablehlo.broadcast_in_dim %36750, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36752 = stablehlo.multiply %36751, %36749 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36753 = stablehlo.multiply %36752, %33214 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36754 = stablehlo.multiply %36753, %33211 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36755 = stablehlo.add %36753, %36754 : tensor<1x1x2048x8192xbf16> loc(#loc4491)
      %36756 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %36757 = stablehlo.broadcast_in_dim %36756, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36758 = stablehlo.multiply %36757, %36755 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36759 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %36760 = stablehlo.broadcast_in_dim %36759, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36761 = stablehlo.multiply %36760, %36758 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36762 = stablehlo.multiply %36761, %33203 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36763 = stablehlo.add %36758, %36762 : tensor<1x1x2048x8192xbf16> loc(#loc4491)
      %36764 = stablehlo.multiply %36748, %33220 : tensor<1x1x2048x8192xbf16> loc(#loc3120)
      %36765 = stablehlo.add %36763, %36764 : tensor<1x1x2048x8192xbf16> loc(#loc4491)
      %36766 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4492)
      %36767 = stablehlo.reduce(%36765 init: %36766) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4492)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4492)
      } loc(#loc4492)
      %36768 = stablehlo.reshape %36767 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4493)
      %36769 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4494)
      %36770 = stablehlo.reduce(%36768 init: %36769) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4494)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4494)
      } loc(#loc4494)
      %36771 = stablehlo.broadcast_in_dim %36770, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4495)
      %36772 = func.call @_where_949(%30664#1, %36771) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %36773 = arith.select %ptrue, %36765, %36765 : tensor<1x1x2048x8192xbf16>
      %36774:2 = func.call @_einsum_950(%30665#0, %33192, %36773) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4496)
      %36775 = func.call @_where_951(%30665#1, %36774#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %36776 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4497)
      %36777 = stablehlo.reduce(%36774#0 init: %36776) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4497)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4497)
      } loc(#loc4497)
      %36778 = stablehlo.reshape %36777 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4498)
      %36779 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4499)
      %36780 = stablehlo.reduce(%36778 init: %36779) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4499)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4499)
      } loc(#loc4499)
      %36781 = stablehlo.broadcast_in_dim %36780, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4500)
      %36782 = func.call @_where_952(%30668#1, %36781) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36783 = stablehlo.multiply %33183, %36774#0 : tensor<1x1x2048x2048xbf16> loc(#loc3113)
      %36784 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4501)
      %36785 = stablehlo.reduce(%36783 init: %36784) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4501)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4501)
      } loc(#loc4501)
      %36786 = stablehlo.reshape %36785 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4502)
      %36787 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4503)
      %36788 = stablehlo.reduce(%36786 init: %36787) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4503)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4503)
      } loc(#loc4503)
      %36789 = stablehlo.broadcast_in_dim %36788, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4504)
      %36790 = func.call @_where_953(%30669#1, %36789) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36791 = stablehlo.broadcast_in_dim %33187, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3113)
      %36792 = stablehlo.multiply %36774#0, %36791 : tensor<1x1x2048x2048xbf16> loc(#loc3113)
      %36793 = stablehlo.multiply %33173, %36792 : tensor<1x1x2048x2048xbf16> loc(#loc3110)
      %36794 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4505)
      %36795 = stablehlo.reduce(%36793 init: %36794) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4505)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4505)
      } loc(#loc4505)
      %36796 = stablehlo.reshape %36795 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4506)
      %36797 = stablehlo.multiply %36796, %33181 : tensor<1x1x2048x1xbf16> loc(#loc3110)
      %36798 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36799 = stablehlo.broadcast_in_dim %36798, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3104)
      %36800 = stablehlo.divide %36797, %36799 : tensor<1x1x2048x1xbf16> loc(#loc3104)
      %36801 = stablehlo.convert %36800 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3100)
      %36802 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4507)
      %36803 = stablehlo.reduce(%36801 init: %36802) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4507)
        stablehlo.return %39096 : tensor<f32> loc(#loc4507)
      } loc(#loc4507)
      %36804 = stablehlo.broadcast_in_dim %36803, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4508)
      %36805 = stablehlo.broadcast_in_dim %36804, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4509)
      %36806 = stablehlo.convert %36805 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3103)
      %36807 = stablehlo.multiply %36806, %33163 : tensor<1x1x2048x2048xbf16> loc(#loc3106)
      %36808 = stablehlo.negate %36807 : tensor<1x1x2048x2048xbf16> loc(#loc4510)
      %36809 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4511)
      %36810 = stablehlo.reduce(%36808 init: %36809) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4511)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4511)
      } loc(#loc4511)
      %36811 = stablehlo.reshape %36810 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4512)
      %36812 = stablehlo.add %36731, %36807 : tensor<1x1x2048x2048xbf16> loc(#loc4513)
      %36813 = stablehlo.broadcast_in_dim %33177, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3110)
      %36814 = stablehlo.multiply %36792, %36813 : tensor<1x1x2048x2048xbf16> loc(#loc3110)
      %36815 = stablehlo.negate %36814 : tensor<1x1x2048x2048xbf16> loc(#loc4510)
      %36816 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4511)
      %36817 = stablehlo.reduce(%36815 init: %36816) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4511)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4511)
      } loc(#loc4511)
      %36818 = stablehlo.reshape %36817 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4512)
      %36819 = stablehlo.add %36812, %36814 : tensor<1x1x2048x2048xbf16> loc(#loc4513)
      %36820 = stablehlo.add %36811, %36818 : tensor<1x1x2048x1xbf16> loc(#loc4513)
      %36821 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36822 = stablehlo.broadcast_in_dim %36821, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3104)
      %36823 = stablehlo.divide %36820, %36822 : tensor<1x1x2048x1xbf16> loc(#loc3104)
      %36824 = stablehlo.convert %36823 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3100)
      %36825 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4507)
      %36826 = stablehlo.reduce(%36824 init: %36825) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4507)
        stablehlo.return %39096 : tensor<f32> loc(#loc4507)
      } loc(#loc4507)
      %36827 = stablehlo.broadcast_in_dim %36826, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4508)
      %36828 = stablehlo.broadcast_in_dim %36827, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4509)
      %36829 = stablehlo.convert %36828 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3103)
      %36830 = stablehlo.add %36819, %36829 : tensor<1x1x2048x2048xbf16> loc(#loc4514)
      %36831 = stablehlo.add %34804, %36830 : tensor<1x1x2048x2048xbf16> loc(#loc4515)
      %36832 = arith.select %ptrue, %36830, %36830 : tensor<1x1x2048x2048xbf16>
      %36833:2 = func.call @_einsum_954(%30674#0, %33145, %36832) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4516)
      %36834 = func.call @_where_955(%30674#1, %36833#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %36835 = arith.select %ptrue, %36833#0, %36833#0 : tensor<1x1x2048x32x64xbf16>
      %36836:2 = func.call @_einsum_956(%33143, %33100, %36835) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4517)
      %36837 = arith.select %ptrue, %36836#1, %36836#1 : tensor<1x1x2048x32x64xbf16>
      %36838 = stablehlo.broadcast_in_dim %36837, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4518)
      %36839 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36840 = stablehlo.pad %36838, %36839, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4519)
      %36841 = stablehlo.convert %36836#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4520)
      %36842 = stablehlo.broadcast_in_dim %33142, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3091)
      %36843 = stablehlo.multiply %36841, %36842 : tensor<1x1x32x2048x2048xf32> loc(#loc3091)
      %36844 = stablehlo.multiply %36843, %33133 : tensor<1x1x32x2048x2048xf32> loc(#loc3091)
      %36845 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3089)
      %36846 = stablehlo.reduce(%36844 init: %36845) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc3089)
        stablehlo.return %39096 : tensor<f32> loc(#loc3089)
      } loc(#loc3089)
      %36847 = stablehlo.reshape %36846 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4521)
      %36848 = stablehlo.negate %36847 : tensor<1x1x32x2048x1xf32> loc(#loc4522)
      %36849 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4523)
      %36850 = stablehlo.reduce(%36848 init: %36849) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4523)
        stablehlo.return %39096 : tensor<f32> loc(#loc4523)
      } loc(#loc4523)
      %36851 = stablehlo.broadcast_in_dim %36850, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4524)
      %36852 = stablehlo.broadcast_in_dim %36851, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4525)
      %36853 = stablehlo.broadcast_in_dim %33136, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc3090)
      %36854 = stablehlo.divide %36841, %36853 : tensor<1x1x32x2048x2048xf32> loc(#loc3090)
      %36855 = stablehlo.add %36852, %36854 : tensor<1x1x32x2048x2048xf32> loc(#loc4526)
      %36856 = stablehlo.multiply %36855, %33133 : tensor<1x1x32x2048x2048xf32> loc(#loc3091)
      %36857 = func.call @_where_957(%33127#1, %36856) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4527)
      %36858 = stablehlo.convert %36857 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4528)
      %36859 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36860 = stablehlo.broadcast_in_dim %36859, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3079)
      %36861 = stablehlo.multiply %36860, %36858 : tensor<1x1x32x2048x2048xbf16> loc(#loc3079)
      %36862 = stablehlo.multiply %36861, %33115 : tensor<1x1x32x2048x2048xbf16> loc(#loc3079)
      %36863 = stablehlo.multiply %36862, %33112 : tensor<1x1x32x2048x2048xbf16> loc(#loc3079)
      %36864 = stablehlo.add %36862, %36863 : tensor<1x1x32x2048x2048xbf16> loc(#loc4529)
      %36865 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %36866 = stablehlo.broadcast_in_dim %36865, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc3076)
      %36867 = stablehlo.divide %36864, %36866 : tensor<1x1x32x2048x2048xbf16> loc(#loc3076)
      %36868:2 = func.call @_einsum_958(%33099, %33107, %36867) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4530)
      %36869 = arith.select %ptrue, %36868#1, %36868#1 : tensor<1x1x2048x32x64xbf16>
      %36870 = stablehlo.broadcast_in_dim %36869, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4518)
      %36871 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36872 = stablehlo.pad %36870, %36871, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4531)
      %36873 = stablehlo.add %36840, %36872 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4532)
      %36874 = stablehlo.multiply %33098, %36868#0 : tensor<1x1x2048x32x64xbf16> loc(#loc3074)
      %36875 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4533)
      %36876 = stablehlo.reduce(%36874 init: %36875) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4533)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4533)
      } loc(#loc4533)
      %36877 = stablehlo.reshape %36876 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4534)
      %36878 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4535)
      %36879 = stablehlo.reduce(%36877 init: %36878) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4535)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4535)
      } loc(#loc4535)
      %36880 = stablehlo.broadcast_in_dim %36879, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4536)
      %36881 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %36882 = stablehlo.broadcast_in_dim %36881, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc3072)
      %36883 = stablehlo.multiply %36882, %36880 : tensor<1x64xbf16> loc(#loc3072)
      %36884 = func.call @softplus_959(%33101#1, %36883) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4537)
      %36885 = func.call @_where_960(%30673#1, %36884) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %36886 = stablehlo.broadcast_in_dim %33105, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc3074)
      %36887 = stablehlo.multiply %36868#0, %36886 : tensor<1x1x2048x32x64xbf16> loc(#loc3074)
      %36888 = arith.select %ptrue, %36887, %36887 : tensor<1x1x2048x32x64xbf16>
      %36889 = stablehlo.broadcast_in_dim %36888, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4518)
      %36890 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %36891 = stablehlo.pad %36889, %36890, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4538)
      %36892 = stablehlo.add %36873, %36891 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4532)
      %36893:2 = func.call @_einsum_961(%33090, %30672#0, %36892) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4539)
      %36894 = func.call @_where_962(%30672#1, %36893#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %36895 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4540)
      %36896 = stablehlo.reduce(%36893#0 init: %36895) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4540)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4540)
      } loc(#loc4540)
      %36897 = stablehlo.reshape %36896 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4541)
      %36898 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4542)
      %36899 = stablehlo.reduce(%36897 init: %36898) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4542)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4542)
      } loc(#loc4542)
      %36900 = stablehlo.broadcast_in_dim %36899, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4543)
      %36901 = func.call @_where_963(%30670#1, %36900) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36902 = stablehlo.multiply %33081, %36893#0 : tensor<1x1x2048x2048xbf16> loc(#loc3062)
      %36903 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4544)
      %36904 = stablehlo.reduce(%36902 init: %36903) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4544)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4544)
      } loc(#loc4544)
      %36905 = stablehlo.reshape %36904 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4545)
      %36906 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4546)
      %36907 = stablehlo.reduce(%36905 init: %36906) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4546)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4546)
      } loc(#loc4546)
      %36908 = stablehlo.broadcast_in_dim %36907, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4547)
      %36909 = func.call @_where_964(%30671#1, %36908) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36910 = stablehlo.broadcast_in_dim %33085, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3062)
      %36911 = stablehlo.multiply %36893#0, %36910 : tensor<1x1x2048x2048xbf16> loc(#loc3062)
      %36912 = stablehlo.multiply %33071, %36911 : tensor<1x1x2048x2048xbf16> loc(#loc3059)
      %36913 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4548)
      %36914 = stablehlo.reduce(%36912 init: %36913) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4548)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4548)
      } loc(#loc4548)
      %36915 = stablehlo.reshape %36914 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4549)
      %36916 = stablehlo.multiply %36915, %33079 : tensor<1x1x2048x1xbf16> loc(#loc3059)
      %36917 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36918 = stablehlo.broadcast_in_dim %36917, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3053)
      %36919 = stablehlo.divide %36916, %36918 : tensor<1x1x2048x1xbf16> loc(#loc3053)
      %36920 = stablehlo.convert %36919 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3049)
      %36921 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4550)
      %36922 = stablehlo.reduce(%36920 init: %36921) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4550)
        stablehlo.return %39096 : tensor<f32> loc(#loc4550)
      } loc(#loc4550)
      %36923 = stablehlo.broadcast_in_dim %36922, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4551)
      %36924 = stablehlo.broadcast_in_dim %36923, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4552)
      %36925 = stablehlo.convert %36924 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3052)
      %36926 = stablehlo.multiply %36925, %33061 : tensor<1x1x2048x2048xbf16> loc(#loc3055)
      %36927 = stablehlo.negate %36926 : tensor<1x1x2048x2048xbf16> loc(#loc4553)
      %36928 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4554)
      %36929 = stablehlo.reduce(%36927 init: %36928) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4554)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4554)
      } loc(#loc4554)
      %36930 = stablehlo.reshape %36929 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4555)
      %36931 = stablehlo.add %36831, %36926 : tensor<1x1x2048x2048xbf16> loc(#loc4556)
      %36932 = stablehlo.broadcast_in_dim %33075, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3059)
      %36933 = stablehlo.multiply %36911, %36932 : tensor<1x1x2048x2048xbf16> loc(#loc3059)
      %36934 = stablehlo.negate %36933 : tensor<1x1x2048x2048xbf16> loc(#loc4553)
      %36935 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4554)
      %36936 = stablehlo.reduce(%36934 init: %36935) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4554)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4554)
      } loc(#loc4554)
      %36937 = stablehlo.reshape %36936 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4555)
      %36938 = stablehlo.add %36931, %36933 : tensor<1x1x2048x2048xbf16> loc(#loc4556)
      %36939 = stablehlo.add %36930, %36937 : tensor<1x1x2048x1xbf16> loc(#loc4556)
      %36940 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %36941 = stablehlo.broadcast_in_dim %36940, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc3053)
      %36942 = stablehlo.divide %36939, %36941 : tensor<1x1x2048x1xbf16> loc(#loc3053)
      %36943 = stablehlo.convert %36942 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc3049)
      %36944 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4550)
      %36945 = stablehlo.reduce(%36943 init: %36944) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4550)
        stablehlo.return %39096 : tensor<f32> loc(#loc4550)
      } loc(#loc4550)
      %36946 = stablehlo.broadcast_in_dim %36945, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4551)
      %36947 = stablehlo.broadcast_in_dim %36946, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4552)
      %36948 = stablehlo.convert %36947 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc3052)
      %36949 = stablehlo.add %36938, %36948 : tensor<1x1x2048x2048xbf16> loc(#loc4557)
      %36950 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %36951 = stablehlo.broadcast_in_dim %36950, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3024)
      %36952 = stablehlo.multiply %36949, %36951 : tensor<1x1x2048x2048xbf16> loc(#loc3024)
      %36953 = stablehlo.broadcast_in_dim %32983, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3023)
      %36954 = stablehlo.multiply %36952, %36953 : tensor<1x1x2048x2048xbf16> loc(#loc3023)
      %36955 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4558)
      %36956 = stablehlo.reduce(%36954 init: %36955) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4558)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4558)
      } loc(#loc4558)
      %36957 = stablehlo.reshape %36956 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4559)
      %36958 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4560)
      %36959 = stablehlo.reduce(%36957 init: %36958) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4560)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4560)
      } loc(#loc4560)
      %36960 = stablehlo.broadcast_in_dim %36959, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4561)
      %36961 = func.call @_where_965(%30655#1, %36960) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %36962 = arith.select %ptrue, %36954, %36954 : tensor<1x1x2048x2048xbf16>
      %36963:2 = func.call @_einsum_966(%30656#0, %32973, %36962) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4562)
      %36964 = func.call @_where_967(%30656#1, %36963#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %36965 = stablehlo.broadcast_in_dim %32971, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3017)
      %36966 = stablehlo.multiply %36963#0, %36965 : tensor<1x1x2048x8192xbf16> loc(#loc3017)
      %36967 = stablehlo.multiply %32942, %36966 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36968 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %36969 = stablehlo.broadcast_in_dim %36968, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36970 = stablehlo.multiply %36969, %36967 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36971 = stablehlo.multiply %36970, %32959 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36972 = stablehlo.multiply %36971, %32956 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36973 = stablehlo.add %36971, %36972 : tensor<1x1x2048x8192xbf16> loc(#loc4563)
      %36974 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %36975 = stablehlo.broadcast_in_dim %36974, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36976 = stablehlo.multiply %36975, %36973 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36977 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %36978 = stablehlo.broadcast_in_dim %36977, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36979 = stablehlo.multiply %36978, %36976 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36980 = stablehlo.multiply %36979, %32948 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36981 = stablehlo.add %36976, %36980 : tensor<1x1x2048x8192xbf16> loc(#loc4563)
      %36982 = stablehlo.multiply %36966, %32965 : tensor<1x1x2048x8192xbf16> loc(#loc3010)
      %36983 = stablehlo.add %36981, %36982 : tensor<1x1x2048x8192xbf16> loc(#loc4563)
      %36984 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4564)
      %36985 = stablehlo.reduce(%36983 init: %36984) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4564)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4564)
      } loc(#loc4564)
      %36986 = stablehlo.reshape %36985 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4565)
      %36987 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4566)
      %36988 = stablehlo.reduce(%36986 init: %36987) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4566)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4566)
      } loc(#loc4566)
      %36989 = stablehlo.broadcast_in_dim %36988, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4567)
      %36990 = func.call @_where_968(%30653#1, %36989) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %36991 = arith.select %ptrue, %36983, %36983 : tensor<1x1x2048x8192xbf16>
      %36992:2 = func.call @_einsum_969(%30654#0, %32937, %36991) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4568)
      %36993 = func.call @_where_970(%30654#1, %36992#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %36994 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4569)
      %36995 = stablehlo.reduce(%36992#0 init: %36994) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4569)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4569)
      } loc(#loc4569)
      %36996 = stablehlo.reshape %36995 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4570)
      %36997 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4571)
      %36998 = stablehlo.reduce(%36996 init: %36997) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4571)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4571)
      } loc(#loc4571)
      %36999 = stablehlo.broadcast_in_dim %36998, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4572)
      %37000 = func.call @_where_971(%30657#1, %36999) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37001 = stablehlo.multiply %32928, %36992#0 : tensor<1x1x2048x2048xbf16> loc(#loc3003)
      %37002 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4573)
      %37003 = stablehlo.reduce(%37001 init: %37002) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4573)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4573)
      } loc(#loc4573)
      %37004 = stablehlo.reshape %37003 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4574)
      %37005 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4575)
      %37006 = stablehlo.reduce(%37004 init: %37005) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4575)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4575)
      } loc(#loc4575)
      %37007 = stablehlo.broadcast_in_dim %37006, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4576)
      %37008 = func.call @_where_972(%30658#1, %37007) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37009 = stablehlo.broadcast_in_dim %32932, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3003)
      %37010 = stablehlo.multiply %36992#0, %37009 : tensor<1x1x2048x2048xbf16> loc(#loc3003)
      %37011 = stablehlo.multiply %32918, %37010 : tensor<1x1x2048x2048xbf16> loc(#loc3000)
      %37012 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4577)
      %37013 = stablehlo.reduce(%37011 init: %37012) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4577)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4577)
      } loc(#loc4577)
      %37014 = stablehlo.reshape %37013 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4578)
      %37015 = stablehlo.multiply %37014, %32926 : tensor<1x1x2048x1xbf16> loc(#loc3000)
      %37016 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37017 = stablehlo.broadcast_in_dim %37016, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2994)
      %37018 = stablehlo.divide %37015, %37017 : tensor<1x1x2048x1xbf16> loc(#loc2994)
      %37019 = stablehlo.convert %37018 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2990)
      %37020 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4579)
      %37021 = stablehlo.reduce(%37019 init: %37020) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4579)
        stablehlo.return %39096 : tensor<f32> loc(#loc4579)
      } loc(#loc4579)
      %37022 = stablehlo.broadcast_in_dim %37021, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4580)
      %37023 = stablehlo.broadcast_in_dim %37022, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4581)
      %37024 = stablehlo.convert %37023 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2993)
      %37025 = stablehlo.multiply %37024, %32908 : tensor<1x1x2048x2048xbf16> loc(#loc2996)
      %37026 = stablehlo.negate %37025 : tensor<1x1x2048x2048xbf16> loc(#loc4582)
      %37027 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4583)
      %37028 = stablehlo.reduce(%37026 init: %37027) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4583)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4583)
      } loc(#loc4583)
      %37029 = stablehlo.reshape %37028 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4584)
      %37030 = stablehlo.add %36949, %37025 : tensor<1x1x2048x2048xbf16> loc(#loc4585)
      %37031 = stablehlo.broadcast_in_dim %32922, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc3000)
      %37032 = stablehlo.multiply %37010, %37031 : tensor<1x1x2048x2048xbf16> loc(#loc3000)
      %37033 = stablehlo.negate %37032 : tensor<1x1x2048x2048xbf16> loc(#loc4582)
      %37034 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4583)
      %37035 = stablehlo.reduce(%37033 init: %37034) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4583)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4583)
      } loc(#loc4583)
      %37036 = stablehlo.reshape %37035 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4584)
      %37037 = stablehlo.add %37030, %37032 : tensor<1x1x2048x2048xbf16> loc(#loc4585)
      %37038 = stablehlo.add %37029, %37036 : tensor<1x1x2048x1xbf16> loc(#loc4585)
      %37039 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37040 = stablehlo.broadcast_in_dim %37039, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2994)
      %37041 = stablehlo.divide %37038, %37040 : tensor<1x1x2048x1xbf16> loc(#loc2994)
      %37042 = stablehlo.convert %37041 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2990)
      %37043 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4579)
      %37044 = stablehlo.reduce(%37042 init: %37043) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4579)
        stablehlo.return %39096 : tensor<f32> loc(#loc4579)
      } loc(#loc4579)
      %37045 = stablehlo.broadcast_in_dim %37044, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4580)
      %37046 = stablehlo.broadcast_in_dim %37045, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4581)
      %37047 = stablehlo.convert %37046 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2993)
      %37048 = stablehlo.add %37037, %37047 : tensor<1x1x2048x2048xbf16> loc(#loc4586)
      %37049 = stablehlo.add %34845, %37048 : tensor<1x1x2048x2048xbf16> loc(#loc4587)
      %37050 = arith.select %ptrue, %37048, %37048 : tensor<1x1x2048x2048xbf16>
      %37051:2 = func.call @_einsum_973(%30663#0, %32890, %37050) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4588)
      %37052 = func.call @_where_974(%30663#1, %37051#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %37053 = arith.select %ptrue, %37051#0, %37051#0 : tensor<1x1x2048x32x64xbf16>
      %37054:2 = func.call @_einsum_975(%32888, %32845, %37053) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4589)
      %37055 = arith.select %ptrue, %37054#1, %37054#1 : tensor<1x1x2048x32x64xbf16>
      %37056 = stablehlo.broadcast_in_dim %37055, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4590)
      %37057 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37058 = stablehlo.pad %37056, %37057, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4591)
      %37059 = stablehlo.convert %37054#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4592)
      %37060 = stablehlo.broadcast_in_dim %32887, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2981)
      %37061 = stablehlo.multiply %37059, %37060 : tensor<1x1x32x2048x2048xf32> loc(#loc2981)
      %37062 = stablehlo.multiply %37061, %32878 : tensor<1x1x32x2048x2048xf32> loc(#loc2981)
      %37063 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2979)
      %37064 = stablehlo.reduce(%37062 init: %37063) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2979)
        stablehlo.return %39096 : tensor<f32> loc(#loc2979)
      } loc(#loc2979)
      %37065 = stablehlo.reshape %37064 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4593)
      %37066 = stablehlo.negate %37065 : tensor<1x1x32x2048x1xf32> loc(#loc4594)
      %37067 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4595)
      %37068 = stablehlo.reduce(%37066 init: %37067) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4595)
        stablehlo.return %39096 : tensor<f32> loc(#loc4595)
      } loc(#loc4595)
      %37069 = stablehlo.broadcast_in_dim %37068, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4596)
      %37070 = stablehlo.broadcast_in_dim %37069, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4597)
      %37071 = stablehlo.broadcast_in_dim %32881, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2980)
      %37072 = stablehlo.divide %37059, %37071 : tensor<1x1x32x2048x2048xf32> loc(#loc2980)
      %37073 = stablehlo.add %37070, %37072 : tensor<1x1x32x2048x2048xf32> loc(#loc4598)
      %37074 = stablehlo.multiply %37073, %32878 : tensor<1x1x32x2048x2048xf32> loc(#loc2981)
      %37075 = func.call @_where_976(%32872#1, %37074) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4599)
      %37076 = stablehlo.convert %37075 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4600)
      %37077 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37078 = stablehlo.broadcast_in_dim %37077, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2969)
      %37079 = stablehlo.multiply %37078, %37076 : tensor<1x1x32x2048x2048xbf16> loc(#loc2969)
      %37080 = stablehlo.multiply %37079, %32860 : tensor<1x1x32x2048x2048xbf16> loc(#loc2969)
      %37081 = stablehlo.multiply %37080, %32857 : tensor<1x1x32x2048x2048xbf16> loc(#loc2969)
      %37082 = stablehlo.add %37080, %37081 : tensor<1x1x32x2048x2048xbf16> loc(#loc4601)
      %37083 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37084 = stablehlo.broadcast_in_dim %37083, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2966)
      %37085 = stablehlo.divide %37082, %37084 : tensor<1x1x32x2048x2048xbf16> loc(#loc2966)
      %37086:2 = func.call @_einsum_977(%32844, %32852, %37085) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4602)
      %37087 = arith.select %ptrue, %37086#1, %37086#1 : tensor<1x1x2048x32x64xbf16>
      %37088 = stablehlo.broadcast_in_dim %37087, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4590)
      %37089 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37090 = stablehlo.pad %37088, %37089, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4603)
      %37091 = stablehlo.add %37058, %37090 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4604)
      %37092 = stablehlo.multiply %32843, %37086#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2964)
      %37093 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4605)
      %37094 = stablehlo.reduce(%37092 init: %37093) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4605)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4605)
      } loc(#loc4605)
      %37095 = stablehlo.reshape %37094 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4606)
      %37096 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4607)
      %37097 = stablehlo.reduce(%37095 init: %37096) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4607)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4607)
      } loc(#loc4607)
      %37098 = stablehlo.broadcast_in_dim %37097, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4608)
      %37099 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %37100 = stablehlo.broadcast_in_dim %37099, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2962)
      %37101 = stablehlo.multiply %37100, %37098 : tensor<1x64xbf16> loc(#loc2962)
      %37102 = func.call @softplus_978(%32846#1, %37101) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4609)
      %37103 = func.call @_where_979(%30662#1, %37102) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %37104 = stablehlo.broadcast_in_dim %32850, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2964)
      %37105 = stablehlo.multiply %37086#0, %37104 : tensor<1x1x2048x32x64xbf16> loc(#loc2964)
      %37106 = arith.select %ptrue, %37105, %37105 : tensor<1x1x2048x32x64xbf16>
      %37107 = stablehlo.broadcast_in_dim %37106, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4590)
      %37108 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37109 = stablehlo.pad %37107, %37108, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4610)
      %37110 = stablehlo.add %37091, %37109 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4604)
      %37111:2 = func.call @_einsum_980(%32835, %30661#0, %37110) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4611)
      %37112 = func.call @_where_981(%30661#1, %37111#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %37113 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4612)
      %37114 = stablehlo.reduce(%37111#0 init: %37113) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4612)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4612)
      } loc(#loc4612)
      %37115 = stablehlo.reshape %37114 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4613)
      %37116 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4614)
      %37117 = stablehlo.reduce(%37115 init: %37116) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4614)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4614)
      } loc(#loc4614)
      %37118 = stablehlo.broadcast_in_dim %37117, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4615)
      %37119 = func.call @_where_982(%30659#1, %37118) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37120 = stablehlo.multiply %32826, %37111#0 : tensor<1x1x2048x2048xbf16> loc(#loc2952)
      %37121 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4616)
      %37122 = stablehlo.reduce(%37120 init: %37121) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4616)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4616)
      } loc(#loc4616)
      %37123 = stablehlo.reshape %37122 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4617)
      %37124 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4618)
      %37125 = stablehlo.reduce(%37123 init: %37124) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4618)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4618)
      } loc(#loc4618)
      %37126 = stablehlo.broadcast_in_dim %37125, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4619)
      %37127 = func.call @_where_983(%30660#1, %37126) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37128 = stablehlo.broadcast_in_dim %32830, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2952)
      %37129 = stablehlo.multiply %37111#0, %37128 : tensor<1x1x2048x2048xbf16> loc(#loc2952)
      %37130 = stablehlo.multiply %32816, %37129 : tensor<1x1x2048x2048xbf16> loc(#loc2949)
      %37131 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4620)
      %37132 = stablehlo.reduce(%37130 init: %37131) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4620)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4620)
      } loc(#loc4620)
      %37133 = stablehlo.reshape %37132 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4621)
      %37134 = stablehlo.multiply %37133, %32824 : tensor<1x1x2048x1xbf16> loc(#loc2949)
      %37135 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37136 = stablehlo.broadcast_in_dim %37135, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2943)
      %37137 = stablehlo.divide %37134, %37136 : tensor<1x1x2048x1xbf16> loc(#loc2943)
      %37138 = stablehlo.convert %37137 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2939)
      %37139 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4622)
      %37140 = stablehlo.reduce(%37138 init: %37139) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4622)
        stablehlo.return %39096 : tensor<f32> loc(#loc4622)
      } loc(#loc4622)
      %37141 = stablehlo.broadcast_in_dim %37140, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4623)
      %37142 = stablehlo.broadcast_in_dim %37141, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4624)
      %37143 = stablehlo.convert %37142 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2942)
      %37144 = stablehlo.multiply %37143, %32806 : tensor<1x1x2048x2048xbf16> loc(#loc2945)
      %37145 = stablehlo.negate %37144 : tensor<1x1x2048x2048xbf16> loc(#loc4625)
      %37146 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4626)
      %37147 = stablehlo.reduce(%37145 init: %37146) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4626)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4626)
      } loc(#loc4626)
      %37148 = stablehlo.reshape %37147 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4627)
      %37149 = stablehlo.add %37049, %37144 : tensor<1x1x2048x2048xbf16> loc(#loc4628)
      %37150 = stablehlo.broadcast_in_dim %32820, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2949)
      %37151 = stablehlo.multiply %37129, %37150 : tensor<1x1x2048x2048xbf16> loc(#loc2949)
      %37152 = stablehlo.negate %37151 : tensor<1x1x2048x2048xbf16> loc(#loc4625)
      %37153 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4626)
      %37154 = stablehlo.reduce(%37152 init: %37153) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4626)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4626)
      } loc(#loc4626)
      %37155 = stablehlo.reshape %37154 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4627)
      %37156 = stablehlo.add %37149, %37151 : tensor<1x1x2048x2048xbf16> loc(#loc4628)
      %37157 = stablehlo.add %37148, %37155 : tensor<1x1x2048x1xbf16> loc(#loc4628)
      %37158 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37159 = stablehlo.broadcast_in_dim %37158, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2943)
      %37160 = stablehlo.divide %37157, %37159 : tensor<1x1x2048x1xbf16> loc(#loc2943)
      %37161 = stablehlo.convert %37160 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2939)
      %37162 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4622)
      %37163 = stablehlo.reduce(%37161 init: %37162) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4622)
        stablehlo.return %39096 : tensor<f32> loc(#loc4622)
      } loc(#loc4622)
      %37164 = stablehlo.broadcast_in_dim %37163, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4623)
      %37165 = stablehlo.broadcast_in_dim %37164, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4624)
      %37166 = stablehlo.convert %37165 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2942)
      %37167 = stablehlo.add %37156, %37166 : tensor<1x1x2048x2048xbf16> loc(#loc4629)
      %37168 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %37169 = stablehlo.broadcast_in_dim %37168, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2914)
      %37170 = stablehlo.multiply %37167, %37169 : tensor<1x1x2048x2048xbf16> loc(#loc2914)
      %37171 = stablehlo.broadcast_in_dim %32728, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2913)
      %37172 = stablehlo.multiply %37170, %37171 : tensor<1x1x2048x2048xbf16> loc(#loc2913)
      %37173 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4630)
      %37174 = stablehlo.reduce(%37172 init: %37173) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4630)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4630)
      } loc(#loc4630)
      %37175 = stablehlo.reshape %37174 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4631)
      %37176 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4632)
      %37177 = stablehlo.reduce(%37175 init: %37176) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4632)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4632)
      } loc(#loc4632)
      %37178 = stablehlo.broadcast_in_dim %37177, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4633)
      %37179 = func.call @_where_984(%30644#1, %37178) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37180 = arith.select %ptrue, %37172, %37172 : tensor<1x1x2048x2048xbf16>
      %37181:2 = func.call @_einsum_985(%30645#0, %32718, %37180) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4634)
      %37182 = func.call @_where_986(%30645#1, %37181#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %37183 = stablehlo.broadcast_in_dim %32716, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2907)
      %37184 = stablehlo.multiply %37181#0, %37183 : tensor<1x1x2048x8192xbf16> loc(#loc2907)
      %37185 = stablehlo.multiply %32687, %37184 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37186 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %37187 = stablehlo.broadcast_in_dim %37186, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37188 = stablehlo.multiply %37187, %37185 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37189 = stablehlo.multiply %37188, %32704 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37190 = stablehlo.multiply %37189, %32701 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37191 = stablehlo.add %37189, %37190 : tensor<1x1x2048x8192xbf16> loc(#loc4635)
      %37192 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %37193 = stablehlo.broadcast_in_dim %37192, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37194 = stablehlo.multiply %37193, %37191 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37195 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %37196 = stablehlo.broadcast_in_dim %37195, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37197 = stablehlo.multiply %37196, %37194 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37198 = stablehlo.multiply %37197, %32693 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37199 = stablehlo.add %37194, %37198 : tensor<1x1x2048x8192xbf16> loc(#loc4635)
      %37200 = stablehlo.multiply %37184, %32710 : tensor<1x1x2048x8192xbf16> loc(#loc2900)
      %37201 = stablehlo.add %37199, %37200 : tensor<1x1x2048x8192xbf16> loc(#loc4635)
      %37202 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4636)
      %37203 = stablehlo.reduce(%37201 init: %37202) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4636)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4636)
      } loc(#loc4636)
      %37204 = stablehlo.reshape %37203 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4637)
      %37205 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4638)
      %37206 = stablehlo.reduce(%37204 init: %37205) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4638)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4638)
      } loc(#loc4638)
      %37207 = stablehlo.broadcast_in_dim %37206, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4639)
      %37208 = func.call @_where_987(%30642#1, %37207) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %37209 = arith.select %ptrue, %37201, %37201 : tensor<1x1x2048x8192xbf16>
      %37210:2 = func.call @_einsum_988(%30643#0, %32682, %37209) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4640)
      %37211 = func.call @_where_989(%30643#1, %37210#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %37212 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4641)
      %37213 = stablehlo.reduce(%37210#0 init: %37212) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4641)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4641)
      } loc(#loc4641)
      %37214 = stablehlo.reshape %37213 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4642)
      %37215 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4643)
      %37216 = stablehlo.reduce(%37214 init: %37215) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4643)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4643)
      } loc(#loc4643)
      %37217 = stablehlo.broadcast_in_dim %37216, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4644)
      %37218 = func.call @_where_990(%30646#1, %37217) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37219 = stablehlo.multiply %32673, %37210#0 : tensor<1x1x2048x2048xbf16> loc(#loc2893)
      %37220 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4645)
      %37221 = stablehlo.reduce(%37219 init: %37220) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4645)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4645)
      } loc(#loc4645)
      %37222 = stablehlo.reshape %37221 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4646)
      %37223 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4647)
      %37224 = stablehlo.reduce(%37222 init: %37223) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4647)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4647)
      } loc(#loc4647)
      %37225 = stablehlo.broadcast_in_dim %37224, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4648)
      %37226 = func.call @_where_991(%30647#1, %37225) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37227 = stablehlo.broadcast_in_dim %32677, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2893)
      %37228 = stablehlo.multiply %37210#0, %37227 : tensor<1x1x2048x2048xbf16> loc(#loc2893)
      %37229 = stablehlo.multiply %32663, %37228 : tensor<1x1x2048x2048xbf16> loc(#loc2890)
      %37230 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4649)
      %37231 = stablehlo.reduce(%37229 init: %37230) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4649)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4649)
      } loc(#loc4649)
      %37232 = stablehlo.reshape %37231 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4650)
      %37233 = stablehlo.multiply %37232, %32671 : tensor<1x1x2048x1xbf16> loc(#loc2890)
      %37234 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37235 = stablehlo.broadcast_in_dim %37234, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2884)
      %37236 = stablehlo.divide %37233, %37235 : tensor<1x1x2048x1xbf16> loc(#loc2884)
      %37237 = stablehlo.convert %37236 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2880)
      %37238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4651)
      %37239 = stablehlo.reduce(%37237 init: %37238) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4651)
        stablehlo.return %39096 : tensor<f32> loc(#loc4651)
      } loc(#loc4651)
      %37240 = stablehlo.broadcast_in_dim %37239, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4652)
      %37241 = stablehlo.broadcast_in_dim %37240, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4653)
      %37242 = stablehlo.convert %37241 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2883)
      %37243 = stablehlo.multiply %37242, %32653 : tensor<1x1x2048x2048xbf16> loc(#loc2886)
      %37244 = stablehlo.negate %37243 : tensor<1x1x2048x2048xbf16> loc(#loc4654)
      %37245 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4655)
      %37246 = stablehlo.reduce(%37244 init: %37245) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4655)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4655)
      } loc(#loc4655)
      %37247 = stablehlo.reshape %37246 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4656)
      %37248 = stablehlo.add %37167, %37243 : tensor<1x1x2048x2048xbf16> loc(#loc4657)
      %37249 = stablehlo.broadcast_in_dim %32667, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2890)
      %37250 = stablehlo.multiply %37228, %37249 : tensor<1x1x2048x2048xbf16> loc(#loc2890)
      %37251 = stablehlo.negate %37250 : tensor<1x1x2048x2048xbf16> loc(#loc4654)
      %37252 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4655)
      %37253 = stablehlo.reduce(%37251 init: %37252) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4655)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4655)
      } loc(#loc4655)
      %37254 = stablehlo.reshape %37253 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4656)
      %37255 = stablehlo.add %37248, %37250 : tensor<1x1x2048x2048xbf16> loc(#loc4657)
      %37256 = stablehlo.add %37247, %37254 : tensor<1x1x2048x1xbf16> loc(#loc4657)
      %37257 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37258 = stablehlo.broadcast_in_dim %37257, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2884)
      %37259 = stablehlo.divide %37256, %37258 : tensor<1x1x2048x1xbf16> loc(#loc2884)
      %37260 = stablehlo.convert %37259 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2880)
      %37261 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4651)
      %37262 = stablehlo.reduce(%37260 init: %37261) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4651)
        stablehlo.return %39096 : tensor<f32> loc(#loc4651)
      } loc(#loc4651)
      %37263 = stablehlo.broadcast_in_dim %37262, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4652)
      %37264 = stablehlo.broadcast_in_dim %37263, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4653)
      %37265 = stablehlo.convert %37264 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2883)
      %37266 = stablehlo.add %37255, %37265 : tensor<1x1x2048x2048xbf16> loc(#loc4658)
      %37267 = stablehlo.add %34886, %37266 : tensor<1x1x2048x2048xbf16> loc(#loc4659)
      %37268 = arith.select %ptrue, %37266, %37266 : tensor<1x1x2048x2048xbf16>
      %37269:2 = func.call @_einsum_992(%30652#0, %32635, %37268) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4660)
      %37270 = func.call @_where_993(%30652#1, %37269#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %37271 = arith.select %ptrue, %37269#0, %37269#0 : tensor<1x1x2048x32x64xbf16>
      %37272:2 = func.call @_einsum_994(%32633, %32590, %37271) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4661)
      %37273 = arith.select %ptrue, %37272#1, %37272#1 : tensor<1x1x2048x32x64xbf16>
      %37274 = stablehlo.broadcast_in_dim %37273, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4662)
      %37275 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37276 = stablehlo.pad %37274, %37275, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4663)
      %37277 = stablehlo.convert %37272#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4664)
      %37278 = stablehlo.broadcast_in_dim %32632, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2871)
      %37279 = stablehlo.multiply %37277, %37278 : tensor<1x1x32x2048x2048xf32> loc(#loc2871)
      %37280 = stablehlo.multiply %37279, %32623 : tensor<1x1x32x2048x2048xf32> loc(#loc2871)
      %37281 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2869)
      %37282 = stablehlo.reduce(%37280 init: %37281) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2869)
        stablehlo.return %39096 : tensor<f32> loc(#loc2869)
      } loc(#loc2869)
      %37283 = stablehlo.reshape %37282 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4665)
      %37284 = stablehlo.negate %37283 : tensor<1x1x32x2048x1xf32> loc(#loc4666)
      %37285 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4667)
      %37286 = stablehlo.reduce(%37284 init: %37285) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4667)
        stablehlo.return %39096 : tensor<f32> loc(#loc4667)
      } loc(#loc4667)
      %37287 = stablehlo.broadcast_in_dim %37286, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4668)
      %37288 = stablehlo.broadcast_in_dim %37287, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4669)
      %37289 = stablehlo.broadcast_in_dim %32626, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2870)
      %37290 = stablehlo.divide %37277, %37289 : tensor<1x1x32x2048x2048xf32> loc(#loc2870)
      %37291 = stablehlo.add %37288, %37290 : tensor<1x1x32x2048x2048xf32> loc(#loc4670)
      %37292 = stablehlo.multiply %37291, %32623 : tensor<1x1x32x2048x2048xf32> loc(#loc2871)
      %37293 = func.call @_where_995(%32617#1, %37292) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4671)
      %37294 = stablehlo.convert %37293 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4672)
      %37295 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37296 = stablehlo.broadcast_in_dim %37295, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2859)
      %37297 = stablehlo.multiply %37296, %37294 : tensor<1x1x32x2048x2048xbf16> loc(#loc2859)
      %37298 = stablehlo.multiply %37297, %32605 : tensor<1x1x32x2048x2048xbf16> loc(#loc2859)
      %37299 = stablehlo.multiply %37298, %32602 : tensor<1x1x32x2048x2048xbf16> loc(#loc2859)
      %37300 = stablehlo.add %37298, %37299 : tensor<1x1x32x2048x2048xbf16> loc(#loc4673)
      %37301 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37302 = stablehlo.broadcast_in_dim %37301, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2856)
      %37303 = stablehlo.divide %37300, %37302 : tensor<1x1x32x2048x2048xbf16> loc(#loc2856)
      %37304:2 = func.call @_einsum_996(%32589, %32597, %37303) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4674)
      %37305 = arith.select %ptrue, %37304#1, %37304#1 : tensor<1x1x2048x32x64xbf16>
      %37306 = stablehlo.broadcast_in_dim %37305, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4662)
      %37307 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37308 = stablehlo.pad %37306, %37307, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4675)
      %37309 = stablehlo.add %37276, %37308 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4676)
      %37310 = stablehlo.multiply %32588, %37304#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2854)
      %37311 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4677)
      %37312 = stablehlo.reduce(%37310 init: %37311) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4677)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4677)
      } loc(#loc4677)
      %37313 = stablehlo.reshape %37312 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4678)
      %37314 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4679)
      %37315 = stablehlo.reduce(%37313 init: %37314) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4679)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4679)
      } loc(#loc4679)
      %37316 = stablehlo.broadcast_in_dim %37315, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4680)
      %37317 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %37318 = stablehlo.broadcast_in_dim %37317, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2852)
      %37319 = stablehlo.multiply %37318, %37316 : tensor<1x64xbf16> loc(#loc2852)
      %37320 = func.call @softplus_997(%32591#1, %37319) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4681)
      %37321 = func.call @_where_998(%30651#1, %37320) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %37322 = stablehlo.broadcast_in_dim %32595, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2854)
      %37323 = stablehlo.multiply %37304#0, %37322 : tensor<1x1x2048x32x64xbf16> loc(#loc2854)
      %37324 = arith.select %ptrue, %37323, %37323 : tensor<1x1x2048x32x64xbf16>
      %37325 = stablehlo.broadcast_in_dim %37324, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4662)
      %37326 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37327 = stablehlo.pad %37325, %37326, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4682)
      %37328 = stablehlo.add %37309, %37327 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4676)
      %37329:2 = func.call @_einsum_999(%32580, %30650#0, %37328) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4683)
      %37330 = func.call @_where_1000(%30650#1, %37329#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %37331 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4684)
      %37332 = stablehlo.reduce(%37329#0 init: %37331) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4684)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4684)
      } loc(#loc4684)
      %37333 = stablehlo.reshape %37332 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4685)
      %37334 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4686)
      %37335 = stablehlo.reduce(%37333 init: %37334) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4686)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4686)
      } loc(#loc4686)
      %37336 = stablehlo.broadcast_in_dim %37335, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4687)
      %37337 = func.call @_where_1001(%30648#1, %37336) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37338 = stablehlo.multiply %32571, %37329#0 : tensor<1x1x2048x2048xbf16> loc(#loc2842)
      %37339 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4688)
      %37340 = stablehlo.reduce(%37338 init: %37339) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4688)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4688)
      } loc(#loc4688)
      %37341 = stablehlo.reshape %37340 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4689)
      %37342 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4690)
      %37343 = stablehlo.reduce(%37341 init: %37342) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4690)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4690)
      } loc(#loc4690)
      %37344 = stablehlo.broadcast_in_dim %37343, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4691)
      %37345 = func.call @_where_1002(%30649#1, %37344) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37346 = stablehlo.broadcast_in_dim %32575, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2842)
      %37347 = stablehlo.multiply %37329#0, %37346 : tensor<1x1x2048x2048xbf16> loc(#loc2842)
      %37348 = stablehlo.multiply %32561, %37347 : tensor<1x1x2048x2048xbf16> loc(#loc2839)
      %37349 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4692)
      %37350 = stablehlo.reduce(%37348 init: %37349) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4692)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4692)
      } loc(#loc4692)
      %37351 = stablehlo.reshape %37350 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4693)
      %37352 = stablehlo.multiply %37351, %32569 : tensor<1x1x2048x1xbf16> loc(#loc2839)
      %37353 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37354 = stablehlo.broadcast_in_dim %37353, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2833)
      %37355 = stablehlo.divide %37352, %37354 : tensor<1x1x2048x1xbf16> loc(#loc2833)
      %37356 = stablehlo.convert %37355 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2829)
      %37357 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4694)
      %37358 = stablehlo.reduce(%37356 init: %37357) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4694)
        stablehlo.return %39096 : tensor<f32> loc(#loc4694)
      } loc(#loc4694)
      %37359 = stablehlo.broadcast_in_dim %37358, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4695)
      %37360 = stablehlo.broadcast_in_dim %37359, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4696)
      %37361 = stablehlo.convert %37360 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2832)
      %37362 = stablehlo.multiply %37361, %32551 : tensor<1x1x2048x2048xbf16> loc(#loc2835)
      %37363 = stablehlo.negate %37362 : tensor<1x1x2048x2048xbf16> loc(#loc4697)
      %37364 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4698)
      %37365 = stablehlo.reduce(%37363 init: %37364) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4698)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4698)
      } loc(#loc4698)
      %37366 = stablehlo.reshape %37365 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4699)
      %37367 = stablehlo.add %37267, %37362 : tensor<1x1x2048x2048xbf16> loc(#loc4700)
      %37368 = stablehlo.broadcast_in_dim %32565, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2839)
      %37369 = stablehlo.multiply %37347, %37368 : tensor<1x1x2048x2048xbf16> loc(#loc2839)
      %37370 = stablehlo.negate %37369 : tensor<1x1x2048x2048xbf16> loc(#loc4697)
      %37371 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4698)
      %37372 = stablehlo.reduce(%37370 init: %37371) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4698)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4698)
      } loc(#loc4698)
      %37373 = stablehlo.reshape %37372 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4699)
      %37374 = stablehlo.add %37367, %37369 : tensor<1x1x2048x2048xbf16> loc(#loc4700)
      %37375 = stablehlo.add %37366, %37373 : tensor<1x1x2048x1xbf16> loc(#loc4700)
      %37376 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37377 = stablehlo.broadcast_in_dim %37376, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2833)
      %37378 = stablehlo.divide %37375, %37377 : tensor<1x1x2048x1xbf16> loc(#loc2833)
      %37379 = stablehlo.convert %37378 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2829)
      %37380 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4694)
      %37381 = stablehlo.reduce(%37379 init: %37380) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4694)
        stablehlo.return %39096 : tensor<f32> loc(#loc4694)
      } loc(#loc4694)
      %37382 = stablehlo.broadcast_in_dim %37381, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4695)
      %37383 = stablehlo.broadcast_in_dim %37382, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4696)
      %37384 = stablehlo.convert %37383 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2832)
      %37385 = stablehlo.add %37374, %37384 : tensor<1x1x2048x2048xbf16> loc(#loc4701)
      %37386 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %37387 = stablehlo.broadcast_in_dim %37386, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2804)
      %37388 = stablehlo.multiply %37385, %37387 : tensor<1x1x2048x2048xbf16> loc(#loc2804)
      %37389 = stablehlo.broadcast_in_dim %32473, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2803)
      %37390 = stablehlo.multiply %37388, %37389 : tensor<1x1x2048x2048xbf16> loc(#loc2803)
      %37391 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4702)
      %37392 = stablehlo.reduce(%37390 init: %37391) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4702)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4702)
      } loc(#loc4702)
      %37393 = stablehlo.reshape %37392 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4703)
      %37394 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4704)
      %37395 = stablehlo.reduce(%37393 init: %37394) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4704)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4704)
      } loc(#loc4704)
      %37396 = stablehlo.broadcast_in_dim %37395, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4705)
      %37397 = func.call @_where_1003(%30633#1, %37396) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37398 = arith.select %ptrue, %37390, %37390 : tensor<1x1x2048x2048xbf16>
      %37399:2 = func.call @_einsum_1004(%30634#0, %32463, %37398) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4706)
      %37400 = func.call @_where_1005(%30634#1, %37399#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %37401 = stablehlo.broadcast_in_dim %32461, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2797)
      %37402 = stablehlo.multiply %37399#0, %37401 : tensor<1x1x2048x8192xbf16> loc(#loc2797)
      %37403 = stablehlo.multiply %32432, %37402 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37404 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %37405 = stablehlo.broadcast_in_dim %37404, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37406 = stablehlo.multiply %37405, %37403 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37407 = stablehlo.multiply %37406, %32449 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37408 = stablehlo.multiply %37407, %32446 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37409 = stablehlo.add %37407, %37408 : tensor<1x1x2048x8192xbf16> loc(#loc4707)
      %37410 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %37411 = stablehlo.broadcast_in_dim %37410, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37412 = stablehlo.multiply %37411, %37409 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37413 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %37414 = stablehlo.broadcast_in_dim %37413, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37415 = stablehlo.multiply %37414, %37412 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37416 = stablehlo.multiply %37415, %32438 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37417 = stablehlo.add %37412, %37416 : tensor<1x1x2048x8192xbf16> loc(#loc4707)
      %37418 = stablehlo.multiply %37402, %32455 : tensor<1x1x2048x8192xbf16> loc(#loc2790)
      %37419 = stablehlo.add %37417, %37418 : tensor<1x1x2048x8192xbf16> loc(#loc4707)
      %37420 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4708)
      %37421 = stablehlo.reduce(%37419 init: %37420) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4708)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4708)
      } loc(#loc4708)
      %37422 = stablehlo.reshape %37421 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4709)
      %37423 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4710)
      %37424 = stablehlo.reduce(%37422 init: %37423) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4710)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4710)
      } loc(#loc4710)
      %37425 = stablehlo.broadcast_in_dim %37424, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4711)
      %37426 = func.call @_where_1006(%30631#1, %37425) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %37427 = arith.select %ptrue, %37419, %37419 : tensor<1x1x2048x8192xbf16>
      %37428:2 = func.call @_einsum_1007(%30632#0, %32427, %37427) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4712)
      %37429 = func.call @_where_1008(%30632#1, %37428#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %37430 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4713)
      %37431 = stablehlo.reduce(%37428#0 init: %37430) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4713)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4713)
      } loc(#loc4713)
      %37432 = stablehlo.reshape %37431 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4714)
      %37433 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4715)
      %37434 = stablehlo.reduce(%37432 init: %37433) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4715)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4715)
      } loc(#loc4715)
      %37435 = stablehlo.broadcast_in_dim %37434, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4716)
      %37436 = func.call @_where_1009(%30635#1, %37435) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37437 = stablehlo.multiply %32418, %37428#0 : tensor<1x1x2048x2048xbf16> loc(#loc2783)
      %37438 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4717)
      %37439 = stablehlo.reduce(%37437 init: %37438) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4717)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4717)
      } loc(#loc4717)
      %37440 = stablehlo.reshape %37439 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4718)
      %37441 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4719)
      %37442 = stablehlo.reduce(%37440 init: %37441) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4719)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4719)
      } loc(#loc4719)
      %37443 = stablehlo.broadcast_in_dim %37442, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4720)
      %37444 = func.call @_where_1010(%30636#1, %37443) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37445 = stablehlo.broadcast_in_dim %32422, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2783)
      %37446 = stablehlo.multiply %37428#0, %37445 : tensor<1x1x2048x2048xbf16> loc(#loc2783)
      %37447 = stablehlo.multiply %32408, %37446 : tensor<1x1x2048x2048xbf16> loc(#loc2780)
      %37448 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4721)
      %37449 = stablehlo.reduce(%37447 init: %37448) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4721)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4721)
      } loc(#loc4721)
      %37450 = stablehlo.reshape %37449 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4722)
      %37451 = stablehlo.multiply %37450, %32416 : tensor<1x1x2048x1xbf16> loc(#loc2780)
      %37452 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37453 = stablehlo.broadcast_in_dim %37452, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2774)
      %37454 = stablehlo.divide %37451, %37453 : tensor<1x1x2048x1xbf16> loc(#loc2774)
      %37455 = stablehlo.convert %37454 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2770)
      %37456 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4723)
      %37457 = stablehlo.reduce(%37455 init: %37456) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4723)
        stablehlo.return %39096 : tensor<f32> loc(#loc4723)
      } loc(#loc4723)
      %37458 = stablehlo.broadcast_in_dim %37457, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4724)
      %37459 = stablehlo.broadcast_in_dim %37458, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4725)
      %37460 = stablehlo.convert %37459 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2773)
      %37461 = stablehlo.multiply %37460, %32398 : tensor<1x1x2048x2048xbf16> loc(#loc2776)
      %37462 = stablehlo.negate %37461 : tensor<1x1x2048x2048xbf16> loc(#loc4726)
      %37463 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4727)
      %37464 = stablehlo.reduce(%37462 init: %37463) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4727)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4727)
      } loc(#loc4727)
      %37465 = stablehlo.reshape %37464 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4728)
      %37466 = stablehlo.add %37385, %37461 : tensor<1x1x2048x2048xbf16> loc(#loc4729)
      %37467 = stablehlo.broadcast_in_dim %32412, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2780)
      %37468 = stablehlo.multiply %37446, %37467 : tensor<1x1x2048x2048xbf16> loc(#loc2780)
      %37469 = stablehlo.negate %37468 : tensor<1x1x2048x2048xbf16> loc(#loc4726)
      %37470 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4727)
      %37471 = stablehlo.reduce(%37469 init: %37470) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4727)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4727)
      } loc(#loc4727)
      %37472 = stablehlo.reshape %37471 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4728)
      %37473 = stablehlo.add %37466, %37468 : tensor<1x1x2048x2048xbf16> loc(#loc4729)
      %37474 = stablehlo.add %37465, %37472 : tensor<1x1x2048x1xbf16> loc(#loc4729)
      %37475 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37476 = stablehlo.broadcast_in_dim %37475, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2774)
      %37477 = stablehlo.divide %37474, %37476 : tensor<1x1x2048x1xbf16> loc(#loc2774)
      %37478 = stablehlo.convert %37477 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2770)
      %37479 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4723)
      %37480 = stablehlo.reduce(%37478 init: %37479) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4723)
        stablehlo.return %39096 : tensor<f32> loc(#loc4723)
      } loc(#loc4723)
      %37481 = stablehlo.broadcast_in_dim %37480, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4724)
      %37482 = stablehlo.broadcast_in_dim %37481, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4725)
      %37483 = stablehlo.convert %37482 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2773)
      %37484 = stablehlo.add %37473, %37483 : tensor<1x1x2048x2048xbf16> loc(#loc4730)
      %37485 = stablehlo.add %34927, %37484 : tensor<1x1x2048x2048xbf16> loc(#loc4731)
      %37486 = arith.select %ptrue, %37484, %37484 : tensor<1x1x2048x2048xbf16>
      %37487:2 = func.call @_einsum_1011(%30641#0, %32380, %37486) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4732)
      %37488 = func.call @_where_1012(%30641#1, %37487#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %37489 = arith.select %ptrue, %37487#0, %37487#0 : tensor<1x1x2048x32x64xbf16>
      %37490:2 = func.call @_einsum_1013(%32378, %32335, %37489) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4733)
      %37491 = arith.select %ptrue, %37490#1, %37490#1 : tensor<1x1x2048x32x64xbf16>
      %37492 = stablehlo.broadcast_in_dim %37491, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4734)
      %37493 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37494 = stablehlo.pad %37492, %37493, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4735)
      %37495 = stablehlo.convert %37490#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4736)
      %37496 = stablehlo.broadcast_in_dim %32377, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2761)
      %37497 = stablehlo.multiply %37495, %37496 : tensor<1x1x32x2048x2048xf32> loc(#loc2761)
      %37498 = stablehlo.multiply %37497, %32368 : tensor<1x1x32x2048x2048xf32> loc(#loc2761)
      %37499 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2759)
      %37500 = stablehlo.reduce(%37498 init: %37499) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2759)
        stablehlo.return %39096 : tensor<f32> loc(#loc2759)
      } loc(#loc2759)
      %37501 = stablehlo.reshape %37500 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4737)
      %37502 = stablehlo.negate %37501 : tensor<1x1x32x2048x1xf32> loc(#loc4738)
      %37503 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4739)
      %37504 = stablehlo.reduce(%37502 init: %37503) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4739)
        stablehlo.return %39096 : tensor<f32> loc(#loc4739)
      } loc(#loc4739)
      %37505 = stablehlo.broadcast_in_dim %37504, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4740)
      %37506 = stablehlo.broadcast_in_dim %37505, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4741)
      %37507 = stablehlo.broadcast_in_dim %32371, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2760)
      %37508 = stablehlo.divide %37495, %37507 : tensor<1x1x32x2048x2048xf32> loc(#loc2760)
      %37509 = stablehlo.add %37506, %37508 : tensor<1x1x32x2048x2048xf32> loc(#loc4742)
      %37510 = stablehlo.multiply %37509, %32368 : tensor<1x1x32x2048x2048xf32> loc(#loc2761)
      %37511 = func.call @_where_1014(%32362#1, %37510) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4743)
      %37512 = stablehlo.convert %37511 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4744)
      %37513 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37514 = stablehlo.broadcast_in_dim %37513, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2749)
      %37515 = stablehlo.multiply %37514, %37512 : tensor<1x1x32x2048x2048xbf16> loc(#loc2749)
      %37516 = stablehlo.multiply %37515, %32350 : tensor<1x1x32x2048x2048xbf16> loc(#loc2749)
      %37517 = stablehlo.multiply %37516, %32347 : tensor<1x1x32x2048x2048xbf16> loc(#loc2749)
      %37518 = stablehlo.add %37516, %37517 : tensor<1x1x32x2048x2048xbf16> loc(#loc4745)
      %37519 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37520 = stablehlo.broadcast_in_dim %37519, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2746)
      %37521 = stablehlo.divide %37518, %37520 : tensor<1x1x32x2048x2048xbf16> loc(#loc2746)
      %37522:2 = func.call @_einsum_1015(%32334, %32342, %37521) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4746)
      %37523 = arith.select %ptrue, %37522#1, %37522#1 : tensor<1x1x2048x32x64xbf16>
      %37524 = stablehlo.broadcast_in_dim %37523, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4734)
      %37525 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37526 = stablehlo.pad %37524, %37525, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4747)
      %37527 = stablehlo.add %37494, %37526 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4748)
      %37528 = stablehlo.multiply %32333, %37522#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2744)
      %37529 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4749)
      %37530 = stablehlo.reduce(%37528 init: %37529) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4749)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4749)
      } loc(#loc4749)
      %37531 = stablehlo.reshape %37530 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4750)
      %37532 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4751)
      %37533 = stablehlo.reduce(%37531 init: %37532) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4751)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4751)
      } loc(#loc4751)
      %37534 = stablehlo.broadcast_in_dim %37533, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4752)
      %37535 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %37536 = stablehlo.broadcast_in_dim %37535, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2742)
      %37537 = stablehlo.multiply %37536, %37534 : tensor<1x64xbf16> loc(#loc2742)
      %37538 = func.call @softplus_1016(%32336#1, %37537) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4753)
      %37539 = func.call @_where_1017(%30640#1, %37538) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %37540 = stablehlo.broadcast_in_dim %32340, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2744)
      %37541 = stablehlo.multiply %37522#0, %37540 : tensor<1x1x2048x32x64xbf16> loc(#loc2744)
      %37542 = arith.select %ptrue, %37541, %37541 : tensor<1x1x2048x32x64xbf16>
      %37543 = stablehlo.broadcast_in_dim %37542, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4734)
      %37544 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37545 = stablehlo.pad %37543, %37544, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4754)
      %37546 = stablehlo.add %37527, %37545 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4748)
      %37547:2 = func.call @_einsum_1018(%32325, %30639#0, %37546) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4755)
      %37548 = func.call @_where_1019(%30639#1, %37547#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %37549 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4756)
      %37550 = stablehlo.reduce(%37547#0 init: %37549) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4756)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4756)
      } loc(#loc4756)
      %37551 = stablehlo.reshape %37550 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4757)
      %37552 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4758)
      %37553 = stablehlo.reduce(%37551 init: %37552) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4758)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4758)
      } loc(#loc4758)
      %37554 = stablehlo.broadcast_in_dim %37553, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4759)
      %37555 = func.call @_where_1020(%30637#1, %37554) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37556 = stablehlo.multiply %32316, %37547#0 : tensor<1x1x2048x2048xbf16> loc(#loc2732)
      %37557 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4760)
      %37558 = stablehlo.reduce(%37556 init: %37557) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4760)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4760)
      } loc(#loc4760)
      %37559 = stablehlo.reshape %37558 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4761)
      %37560 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4762)
      %37561 = stablehlo.reduce(%37559 init: %37560) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4762)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4762)
      } loc(#loc4762)
      %37562 = stablehlo.broadcast_in_dim %37561, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4763)
      %37563 = func.call @_where_1021(%30638#1, %37562) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37564 = stablehlo.broadcast_in_dim %32320, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2732)
      %37565 = stablehlo.multiply %37547#0, %37564 : tensor<1x1x2048x2048xbf16> loc(#loc2732)
      %37566 = stablehlo.multiply %32306, %37565 : tensor<1x1x2048x2048xbf16> loc(#loc2729)
      %37567 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4764)
      %37568 = stablehlo.reduce(%37566 init: %37567) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4764)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4764)
      } loc(#loc4764)
      %37569 = stablehlo.reshape %37568 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4765)
      %37570 = stablehlo.multiply %37569, %32314 : tensor<1x1x2048x1xbf16> loc(#loc2729)
      %37571 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37572 = stablehlo.broadcast_in_dim %37571, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2723)
      %37573 = stablehlo.divide %37570, %37572 : tensor<1x1x2048x1xbf16> loc(#loc2723)
      %37574 = stablehlo.convert %37573 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2719)
      %37575 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4766)
      %37576 = stablehlo.reduce(%37574 init: %37575) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4766)
        stablehlo.return %39096 : tensor<f32> loc(#loc4766)
      } loc(#loc4766)
      %37577 = stablehlo.broadcast_in_dim %37576, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4767)
      %37578 = stablehlo.broadcast_in_dim %37577, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4768)
      %37579 = stablehlo.convert %37578 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2722)
      %37580 = stablehlo.multiply %37579, %32296 : tensor<1x1x2048x2048xbf16> loc(#loc2725)
      %37581 = stablehlo.negate %37580 : tensor<1x1x2048x2048xbf16> loc(#loc4769)
      %37582 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4770)
      %37583 = stablehlo.reduce(%37581 init: %37582) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4770)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4770)
      } loc(#loc4770)
      %37584 = stablehlo.reshape %37583 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4771)
      %37585 = stablehlo.add %37485, %37580 : tensor<1x1x2048x2048xbf16> loc(#loc4772)
      %37586 = stablehlo.broadcast_in_dim %32310, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2729)
      %37587 = stablehlo.multiply %37565, %37586 : tensor<1x1x2048x2048xbf16> loc(#loc2729)
      %37588 = stablehlo.negate %37587 : tensor<1x1x2048x2048xbf16> loc(#loc4769)
      %37589 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4770)
      %37590 = stablehlo.reduce(%37588 init: %37589) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4770)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4770)
      } loc(#loc4770)
      %37591 = stablehlo.reshape %37590 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4771)
      %37592 = stablehlo.add %37585, %37587 : tensor<1x1x2048x2048xbf16> loc(#loc4772)
      %37593 = stablehlo.add %37584, %37591 : tensor<1x1x2048x1xbf16> loc(#loc4772)
      %37594 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37595 = stablehlo.broadcast_in_dim %37594, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2723)
      %37596 = stablehlo.divide %37593, %37595 : tensor<1x1x2048x1xbf16> loc(#loc2723)
      %37597 = stablehlo.convert %37596 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2719)
      %37598 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4766)
      %37599 = stablehlo.reduce(%37597 init: %37598) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4766)
        stablehlo.return %39096 : tensor<f32> loc(#loc4766)
      } loc(#loc4766)
      %37600 = stablehlo.broadcast_in_dim %37599, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4767)
      %37601 = stablehlo.broadcast_in_dim %37600, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4768)
      %37602 = stablehlo.convert %37601 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2722)
      %37603 = stablehlo.add %37592, %37602 : tensor<1x1x2048x2048xbf16> loc(#loc4773)
      %37604 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %37605 = stablehlo.broadcast_in_dim %37604, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2694)
      %37606 = stablehlo.multiply %37603, %37605 : tensor<1x1x2048x2048xbf16> loc(#loc2694)
      %37607 = stablehlo.broadcast_in_dim %32218, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2693)
      %37608 = stablehlo.multiply %37606, %37607 : tensor<1x1x2048x2048xbf16> loc(#loc2693)
      %37609 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4774)
      %37610 = stablehlo.reduce(%37608 init: %37609) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4774)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4774)
      } loc(#loc4774)
      %37611 = stablehlo.reshape %37610 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4775)
      %37612 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4776)
      %37613 = stablehlo.reduce(%37611 init: %37612) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4776)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4776)
      } loc(#loc4776)
      %37614 = stablehlo.broadcast_in_dim %37613, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4777)
      %37615 = func.call @_where_1022(%30622#1, %37614) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37616 = arith.select %ptrue, %37608, %37608 : tensor<1x1x2048x2048xbf16>
      %37617:2 = func.call @_einsum_1023(%30623#0, %32208, %37616) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4778)
      %37618 = func.call @_where_1024(%30623#1, %37617#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %37619 = stablehlo.broadcast_in_dim %32206, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2687)
      %37620 = stablehlo.multiply %37617#0, %37619 : tensor<1x1x2048x8192xbf16> loc(#loc2687)
      %37621 = stablehlo.multiply %32177, %37620 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37622 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %37623 = stablehlo.broadcast_in_dim %37622, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37624 = stablehlo.multiply %37623, %37621 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37625 = stablehlo.multiply %37624, %32194 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37626 = stablehlo.multiply %37625, %32191 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37627 = stablehlo.add %37625, %37626 : tensor<1x1x2048x8192xbf16> loc(#loc4779)
      %37628 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %37629 = stablehlo.broadcast_in_dim %37628, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37630 = stablehlo.multiply %37629, %37627 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37631 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %37632 = stablehlo.broadcast_in_dim %37631, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37633 = stablehlo.multiply %37632, %37630 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37634 = stablehlo.multiply %37633, %32183 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37635 = stablehlo.add %37630, %37634 : tensor<1x1x2048x8192xbf16> loc(#loc4779)
      %37636 = stablehlo.multiply %37620, %32200 : tensor<1x1x2048x8192xbf16> loc(#loc2680)
      %37637 = stablehlo.add %37635, %37636 : tensor<1x1x2048x8192xbf16> loc(#loc4779)
      %37638 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4780)
      %37639 = stablehlo.reduce(%37637 init: %37638) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4780)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4780)
      } loc(#loc4780)
      %37640 = stablehlo.reshape %37639 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4781)
      %37641 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4782)
      %37642 = stablehlo.reduce(%37640 init: %37641) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4782)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4782)
      } loc(#loc4782)
      %37643 = stablehlo.broadcast_in_dim %37642, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4783)
      %37644 = func.call @_where_1025(%30620#1, %37643) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %37645 = arith.select %ptrue, %37637, %37637 : tensor<1x1x2048x8192xbf16>
      %37646:2 = func.call @_einsum_1026(%30621#0, %32172, %37645) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4784)
      %37647 = func.call @_where_1027(%30621#1, %37646#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %37648 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4785)
      %37649 = stablehlo.reduce(%37646#0 init: %37648) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4785)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4785)
      } loc(#loc4785)
      %37650 = stablehlo.reshape %37649 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4786)
      %37651 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4787)
      %37652 = stablehlo.reduce(%37650 init: %37651) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4787)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4787)
      } loc(#loc4787)
      %37653 = stablehlo.broadcast_in_dim %37652, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4788)
      %37654 = func.call @_where_1028(%30624#1, %37653) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37655 = stablehlo.multiply %32163, %37646#0 : tensor<1x1x2048x2048xbf16> loc(#loc2673)
      %37656 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4789)
      %37657 = stablehlo.reduce(%37655 init: %37656) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4789)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4789)
      } loc(#loc4789)
      %37658 = stablehlo.reshape %37657 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4790)
      %37659 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4791)
      %37660 = stablehlo.reduce(%37658 init: %37659) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4791)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4791)
      } loc(#loc4791)
      %37661 = stablehlo.broadcast_in_dim %37660, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4792)
      %37662 = func.call @_where_1029(%30625#1, %37661) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37663 = stablehlo.broadcast_in_dim %32167, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2673)
      %37664 = stablehlo.multiply %37646#0, %37663 : tensor<1x1x2048x2048xbf16> loc(#loc2673)
      %37665 = stablehlo.multiply %32153, %37664 : tensor<1x1x2048x2048xbf16> loc(#loc2670)
      %37666 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4793)
      %37667 = stablehlo.reduce(%37665 init: %37666) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4793)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4793)
      } loc(#loc4793)
      %37668 = stablehlo.reshape %37667 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4794)
      %37669 = stablehlo.multiply %37668, %32161 : tensor<1x1x2048x1xbf16> loc(#loc2670)
      %37670 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37671 = stablehlo.broadcast_in_dim %37670, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2664)
      %37672 = stablehlo.divide %37669, %37671 : tensor<1x1x2048x1xbf16> loc(#loc2664)
      %37673 = stablehlo.convert %37672 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2660)
      %37674 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4795)
      %37675 = stablehlo.reduce(%37673 init: %37674) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4795)
        stablehlo.return %39096 : tensor<f32> loc(#loc4795)
      } loc(#loc4795)
      %37676 = stablehlo.broadcast_in_dim %37675, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4796)
      %37677 = stablehlo.broadcast_in_dim %37676, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4797)
      %37678 = stablehlo.convert %37677 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2663)
      %37679 = stablehlo.multiply %37678, %32143 : tensor<1x1x2048x2048xbf16> loc(#loc2666)
      %37680 = stablehlo.negate %37679 : tensor<1x1x2048x2048xbf16> loc(#loc4798)
      %37681 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4799)
      %37682 = stablehlo.reduce(%37680 init: %37681) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4799)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4799)
      } loc(#loc4799)
      %37683 = stablehlo.reshape %37682 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4800)
      %37684 = stablehlo.add %37603, %37679 : tensor<1x1x2048x2048xbf16> loc(#loc4801)
      %37685 = stablehlo.broadcast_in_dim %32157, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2670)
      %37686 = stablehlo.multiply %37664, %37685 : tensor<1x1x2048x2048xbf16> loc(#loc2670)
      %37687 = stablehlo.negate %37686 : tensor<1x1x2048x2048xbf16> loc(#loc4798)
      %37688 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4799)
      %37689 = stablehlo.reduce(%37687 init: %37688) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4799)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4799)
      } loc(#loc4799)
      %37690 = stablehlo.reshape %37689 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4800)
      %37691 = stablehlo.add %37684, %37686 : tensor<1x1x2048x2048xbf16> loc(#loc4801)
      %37692 = stablehlo.add %37683, %37690 : tensor<1x1x2048x1xbf16> loc(#loc4801)
      %37693 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37694 = stablehlo.broadcast_in_dim %37693, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2664)
      %37695 = stablehlo.divide %37692, %37694 : tensor<1x1x2048x1xbf16> loc(#loc2664)
      %37696 = stablehlo.convert %37695 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2660)
      %37697 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4795)
      %37698 = stablehlo.reduce(%37696 init: %37697) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4795)
        stablehlo.return %39096 : tensor<f32> loc(#loc4795)
      } loc(#loc4795)
      %37699 = stablehlo.broadcast_in_dim %37698, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4796)
      %37700 = stablehlo.broadcast_in_dim %37699, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4797)
      %37701 = stablehlo.convert %37700 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2663)
      %37702 = stablehlo.add %37691, %37701 : tensor<1x1x2048x2048xbf16> loc(#loc4802)
      %37703 = stablehlo.add %34968, %37702 : tensor<1x1x2048x2048xbf16> loc(#loc4803)
      %37704 = arith.select %ptrue, %37702, %37702 : tensor<1x1x2048x2048xbf16>
      %37705:2 = func.call @_einsum_1030(%30630#0, %32125, %37704) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4804)
      %37706 = func.call @_where_1031(%30630#1, %37705#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %37707 = arith.select %ptrue, %37705#0, %37705#0 : tensor<1x1x2048x32x64xbf16>
      %37708:2 = func.call @_einsum_1032(%32123, %32080, %37707) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4805)
      %37709 = arith.select %ptrue, %37708#1, %37708#1 : tensor<1x1x2048x32x64xbf16>
      %37710 = stablehlo.broadcast_in_dim %37709, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4806)
      %37711 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37712 = stablehlo.pad %37710, %37711, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4807)
      %37713 = stablehlo.convert %37708#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4808)
      %37714 = stablehlo.broadcast_in_dim %32122, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2651)
      %37715 = stablehlo.multiply %37713, %37714 : tensor<1x1x32x2048x2048xf32> loc(#loc2651)
      %37716 = stablehlo.multiply %37715, %32113 : tensor<1x1x32x2048x2048xf32> loc(#loc2651)
      %37717 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2649)
      %37718 = stablehlo.reduce(%37716 init: %37717) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2649)
        stablehlo.return %39096 : tensor<f32> loc(#loc2649)
      } loc(#loc2649)
      %37719 = stablehlo.reshape %37718 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4809)
      %37720 = stablehlo.negate %37719 : tensor<1x1x32x2048x1xf32> loc(#loc4810)
      %37721 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4811)
      %37722 = stablehlo.reduce(%37720 init: %37721) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4811)
        stablehlo.return %39096 : tensor<f32> loc(#loc4811)
      } loc(#loc4811)
      %37723 = stablehlo.broadcast_in_dim %37722, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4812)
      %37724 = stablehlo.broadcast_in_dim %37723, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4813)
      %37725 = stablehlo.broadcast_in_dim %32116, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2650)
      %37726 = stablehlo.divide %37713, %37725 : tensor<1x1x32x2048x2048xf32> loc(#loc2650)
      %37727 = stablehlo.add %37724, %37726 : tensor<1x1x32x2048x2048xf32> loc(#loc4814)
      %37728 = stablehlo.multiply %37727, %32113 : tensor<1x1x32x2048x2048xf32> loc(#loc2651)
      %37729 = func.call @_where_1033(%32107#1, %37728) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4815)
      %37730 = stablehlo.convert %37729 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4816)
      %37731 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37732 = stablehlo.broadcast_in_dim %37731, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2639)
      %37733 = stablehlo.multiply %37732, %37730 : tensor<1x1x32x2048x2048xbf16> loc(#loc2639)
      %37734 = stablehlo.multiply %37733, %32095 : tensor<1x1x32x2048x2048xbf16> loc(#loc2639)
      %37735 = stablehlo.multiply %37734, %32092 : tensor<1x1x32x2048x2048xbf16> loc(#loc2639)
      %37736 = stablehlo.add %37734, %37735 : tensor<1x1x32x2048x2048xbf16> loc(#loc4817)
      %37737 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37738 = stablehlo.broadcast_in_dim %37737, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2636)
      %37739 = stablehlo.divide %37736, %37738 : tensor<1x1x32x2048x2048xbf16> loc(#loc2636)
      %37740:2 = func.call @_einsum_1034(%32079, %32087, %37739) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4818)
      %37741 = arith.select %ptrue, %37740#1, %37740#1 : tensor<1x1x2048x32x64xbf16>
      %37742 = stablehlo.broadcast_in_dim %37741, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4806)
      %37743 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37744 = stablehlo.pad %37742, %37743, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4819)
      %37745 = stablehlo.add %37712, %37744 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4820)
      %37746 = stablehlo.multiply %32078, %37740#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2634)
      %37747 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4821)
      %37748 = stablehlo.reduce(%37746 init: %37747) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4821)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4821)
      } loc(#loc4821)
      %37749 = stablehlo.reshape %37748 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4822)
      %37750 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4823)
      %37751 = stablehlo.reduce(%37749 init: %37750) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4823)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4823)
      } loc(#loc4823)
      %37752 = stablehlo.broadcast_in_dim %37751, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4824)
      %37753 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %37754 = stablehlo.broadcast_in_dim %37753, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2632)
      %37755 = stablehlo.multiply %37754, %37752 : tensor<1x64xbf16> loc(#loc2632)
      %37756 = func.call @softplus_1035(%32081#1, %37755) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4825)
      %37757 = func.call @_where_1036(%30629#1, %37756) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %37758 = stablehlo.broadcast_in_dim %32085, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2634)
      %37759 = stablehlo.multiply %37740#0, %37758 : tensor<1x1x2048x32x64xbf16> loc(#loc2634)
      %37760 = arith.select %ptrue, %37759, %37759 : tensor<1x1x2048x32x64xbf16>
      %37761 = stablehlo.broadcast_in_dim %37760, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4806)
      %37762 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37763 = stablehlo.pad %37761, %37762, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4826)
      %37764 = stablehlo.add %37745, %37763 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4820)
      %37765:2 = func.call @_einsum_1037(%32070, %30628#0, %37764) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4827)
      %37766 = func.call @_where_1038(%30628#1, %37765#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %37767 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4828)
      %37768 = stablehlo.reduce(%37765#0 init: %37767) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4828)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4828)
      } loc(#loc4828)
      %37769 = stablehlo.reshape %37768 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4829)
      %37770 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4830)
      %37771 = stablehlo.reduce(%37769 init: %37770) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4830)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4830)
      } loc(#loc4830)
      %37772 = stablehlo.broadcast_in_dim %37771, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4831)
      %37773 = func.call @_where_1039(%30626#1, %37772) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37774 = stablehlo.multiply %32061, %37765#0 : tensor<1x1x2048x2048xbf16> loc(#loc2622)
      %37775 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4832)
      %37776 = stablehlo.reduce(%37774 init: %37775) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4832)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4832)
      } loc(#loc4832)
      %37777 = stablehlo.reshape %37776 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4833)
      %37778 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4834)
      %37779 = stablehlo.reduce(%37777 init: %37778) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4834)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4834)
      } loc(#loc4834)
      %37780 = stablehlo.broadcast_in_dim %37779, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4835)
      %37781 = func.call @_where_1040(%30627#1, %37780) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37782 = stablehlo.broadcast_in_dim %32065, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2622)
      %37783 = stablehlo.multiply %37765#0, %37782 : tensor<1x1x2048x2048xbf16> loc(#loc2622)
      %37784 = stablehlo.multiply %32051, %37783 : tensor<1x1x2048x2048xbf16> loc(#loc2619)
      %37785 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4836)
      %37786 = stablehlo.reduce(%37784 init: %37785) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4836)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4836)
      } loc(#loc4836)
      %37787 = stablehlo.reshape %37786 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4837)
      %37788 = stablehlo.multiply %37787, %32059 : tensor<1x1x2048x1xbf16> loc(#loc2619)
      %37789 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37790 = stablehlo.broadcast_in_dim %37789, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2613)
      %37791 = stablehlo.divide %37788, %37790 : tensor<1x1x2048x1xbf16> loc(#loc2613)
      %37792 = stablehlo.convert %37791 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2609)
      %37793 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4838)
      %37794 = stablehlo.reduce(%37792 init: %37793) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4838)
        stablehlo.return %39096 : tensor<f32> loc(#loc4838)
      } loc(#loc4838)
      %37795 = stablehlo.broadcast_in_dim %37794, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4839)
      %37796 = stablehlo.broadcast_in_dim %37795, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4840)
      %37797 = stablehlo.convert %37796 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2612)
      %37798 = stablehlo.multiply %37797, %32041 : tensor<1x1x2048x2048xbf16> loc(#loc2615)
      %37799 = stablehlo.negate %37798 : tensor<1x1x2048x2048xbf16> loc(#loc4841)
      %37800 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4842)
      %37801 = stablehlo.reduce(%37799 init: %37800) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4842)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4842)
      } loc(#loc4842)
      %37802 = stablehlo.reshape %37801 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4843)
      %37803 = stablehlo.add %37703, %37798 : tensor<1x1x2048x2048xbf16> loc(#loc4844)
      %37804 = stablehlo.broadcast_in_dim %32055, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2619)
      %37805 = stablehlo.multiply %37783, %37804 : tensor<1x1x2048x2048xbf16> loc(#loc2619)
      %37806 = stablehlo.negate %37805 : tensor<1x1x2048x2048xbf16> loc(#loc4841)
      %37807 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4842)
      %37808 = stablehlo.reduce(%37806 init: %37807) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4842)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4842)
      } loc(#loc4842)
      %37809 = stablehlo.reshape %37808 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4843)
      %37810 = stablehlo.add %37803, %37805 : tensor<1x1x2048x2048xbf16> loc(#loc4844)
      %37811 = stablehlo.add %37802, %37809 : tensor<1x1x2048x1xbf16> loc(#loc4844)
      %37812 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37813 = stablehlo.broadcast_in_dim %37812, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2613)
      %37814 = stablehlo.divide %37811, %37813 : tensor<1x1x2048x1xbf16> loc(#loc2613)
      %37815 = stablehlo.convert %37814 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2609)
      %37816 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4838)
      %37817 = stablehlo.reduce(%37815 init: %37816) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4838)
        stablehlo.return %39096 : tensor<f32> loc(#loc4838)
      } loc(#loc4838)
      %37818 = stablehlo.broadcast_in_dim %37817, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4839)
      %37819 = stablehlo.broadcast_in_dim %37818, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4840)
      %37820 = stablehlo.convert %37819 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2612)
      %37821 = stablehlo.add %37810, %37820 : tensor<1x1x2048x2048xbf16> loc(#loc4845)
      %37822 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %37823 = stablehlo.broadcast_in_dim %37822, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2584)
      %37824 = stablehlo.multiply %37821, %37823 : tensor<1x1x2048x2048xbf16> loc(#loc2584)
      %37825 = stablehlo.broadcast_in_dim %31963, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2583)
      %37826 = stablehlo.multiply %37824, %37825 : tensor<1x1x2048x2048xbf16> loc(#loc2583)
      %37827 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4846)
      %37828 = stablehlo.reduce(%37826 init: %37827) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4846)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4846)
      } loc(#loc4846)
      %37829 = stablehlo.reshape %37828 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4847)
      %37830 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4848)
      %37831 = stablehlo.reduce(%37829 init: %37830) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4848)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4848)
      } loc(#loc4848)
      %37832 = stablehlo.broadcast_in_dim %37831, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4849)
      %37833 = func.call @_where_1041(%30611#1, %37832) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37834 = arith.select %ptrue, %37826, %37826 : tensor<1x1x2048x2048xbf16>
      %37835:2 = func.call @_einsum_1042(%30612#0, %31953, %37834) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4850)
      %37836 = func.call @_where_1043(%30612#1, %37835#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %37837 = stablehlo.broadcast_in_dim %31951, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2577)
      %37838 = stablehlo.multiply %37835#0, %37837 : tensor<1x1x2048x8192xbf16> loc(#loc2577)
      %37839 = stablehlo.multiply %31922, %37838 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37840 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %37841 = stablehlo.broadcast_in_dim %37840, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37842 = stablehlo.multiply %37841, %37839 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37843 = stablehlo.multiply %37842, %31939 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37844 = stablehlo.multiply %37843, %31936 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37845 = stablehlo.add %37843, %37844 : tensor<1x1x2048x8192xbf16> loc(#loc4851)
      %37846 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %37847 = stablehlo.broadcast_in_dim %37846, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37848 = stablehlo.multiply %37847, %37845 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37849 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %37850 = stablehlo.broadcast_in_dim %37849, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37851 = stablehlo.multiply %37850, %37848 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37852 = stablehlo.multiply %37851, %31928 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37853 = stablehlo.add %37848, %37852 : tensor<1x1x2048x8192xbf16> loc(#loc4851)
      %37854 = stablehlo.multiply %37838, %31945 : tensor<1x1x2048x8192xbf16> loc(#loc2570)
      %37855 = stablehlo.add %37853, %37854 : tensor<1x1x2048x8192xbf16> loc(#loc4851)
      %37856 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4852)
      %37857 = stablehlo.reduce(%37855 init: %37856) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4852)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4852)
      } loc(#loc4852)
      %37858 = stablehlo.reshape %37857 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4853)
      %37859 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4854)
      %37860 = stablehlo.reduce(%37858 init: %37859) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4854)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4854)
      } loc(#loc4854)
      %37861 = stablehlo.broadcast_in_dim %37860, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4855)
      %37862 = func.call @_where_1044(%30609#1, %37861) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %37863 = arith.select %ptrue, %37855, %37855 : tensor<1x1x2048x8192xbf16>
      %37864:2 = func.call @_einsum_1045(%30610#0, %31917, %37863) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4856)
      %37865 = func.call @_where_1046(%30610#1, %37864#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %37866 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4857)
      %37867 = stablehlo.reduce(%37864#0 init: %37866) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4857)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4857)
      } loc(#loc4857)
      %37868 = stablehlo.reshape %37867 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4858)
      %37869 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4859)
      %37870 = stablehlo.reduce(%37868 init: %37869) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4859)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4859)
      } loc(#loc4859)
      %37871 = stablehlo.broadcast_in_dim %37870, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4860)
      %37872 = func.call @_where_1047(%30613#1, %37871) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37873 = stablehlo.multiply %31908, %37864#0 : tensor<1x1x2048x2048xbf16> loc(#loc2563)
      %37874 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4861)
      %37875 = stablehlo.reduce(%37873 init: %37874) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4861)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4861)
      } loc(#loc4861)
      %37876 = stablehlo.reshape %37875 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4862)
      %37877 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4863)
      %37878 = stablehlo.reduce(%37876 init: %37877) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4863)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4863)
      } loc(#loc4863)
      %37879 = stablehlo.broadcast_in_dim %37878, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4864)
      %37880 = func.call @_where_1048(%30614#1, %37879) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37881 = stablehlo.broadcast_in_dim %31912, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2563)
      %37882 = stablehlo.multiply %37864#0, %37881 : tensor<1x1x2048x2048xbf16> loc(#loc2563)
      %37883 = stablehlo.multiply %31898, %37882 : tensor<1x1x2048x2048xbf16> loc(#loc2560)
      %37884 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4865)
      %37885 = stablehlo.reduce(%37883 init: %37884) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4865)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4865)
      } loc(#loc4865)
      %37886 = stablehlo.reshape %37885 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4866)
      %37887 = stablehlo.multiply %37886, %31906 : tensor<1x1x2048x1xbf16> loc(#loc2560)
      %37888 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37889 = stablehlo.broadcast_in_dim %37888, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2554)
      %37890 = stablehlo.divide %37887, %37889 : tensor<1x1x2048x1xbf16> loc(#loc2554)
      %37891 = stablehlo.convert %37890 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2550)
      %37892 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4867)
      %37893 = stablehlo.reduce(%37891 init: %37892) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4867)
        stablehlo.return %39096 : tensor<f32> loc(#loc4867)
      } loc(#loc4867)
      %37894 = stablehlo.broadcast_in_dim %37893, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4868)
      %37895 = stablehlo.broadcast_in_dim %37894, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4869)
      %37896 = stablehlo.convert %37895 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2553)
      %37897 = stablehlo.multiply %37896, %31888 : tensor<1x1x2048x2048xbf16> loc(#loc2556)
      %37898 = stablehlo.negate %37897 : tensor<1x1x2048x2048xbf16> loc(#loc4870)
      %37899 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4871)
      %37900 = stablehlo.reduce(%37898 init: %37899) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4871)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4871)
      } loc(#loc4871)
      %37901 = stablehlo.reshape %37900 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4872)
      %37902 = stablehlo.add %37821, %37897 : tensor<1x1x2048x2048xbf16> loc(#loc4873)
      %37903 = stablehlo.broadcast_in_dim %31902, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2560)
      %37904 = stablehlo.multiply %37882, %37903 : tensor<1x1x2048x2048xbf16> loc(#loc2560)
      %37905 = stablehlo.negate %37904 : tensor<1x1x2048x2048xbf16> loc(#loc4870)
      %37906 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4871)
      %37907 = stablehlo.reduce(%37905 init: %37906) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4871)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4871)
      } loc(#loc4871)
      %37908 = stablehlo.reshape %37907 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4872)
      %37909 = stablehlo.add %37902, %37904 : tensor<1x1x2048x2048xbf16> loc(#loc4873)
      %37910 = stablehlo.add %37901, %37908 : tensor<1x1x2048x1xbf16> loc(#loc4873)
      %37911 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %37912 = stablehlo.broadcast_in_dim %37911, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2554)
      %37913 = stablehlo.divide %37910, %37912 : tensor<1x1x2048x1xbf16> loc(#loc2554)
      %37914 = stablehlo.convert %37913 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2550)
      %37915 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4867)
      %37916 = stablehlo.reduce(%37914 init: %37915) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4867)
        stablehlo.return %39096 : tensor<f32> loc(#loc4867)
      } loc(#loc4867)
      %37917 = stablehlo.broadcast_in_dim %37916, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4868)
      %37918 = stablehlo.broadcast_in_dim %37917, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4869)
      %37919 = stablehlo.convert %37918 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2553)
      %37920 = stablehlo.add %37909, %37919 : tensor<1x1x2048x2048xbf16> loc(#loc4874)
      %37921 = stablehlo.add %35009, %37920 : tensor<1x1x2048x2048xbf16> loc(#loc4875)
      %37922 = arith.select %ptrue, %37920, %37920 : tensor<1x1x2048x2048xbf16>
      %37923:2 = func.call @_einsum_1049(%30619#0, %31870, %37922) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4876)
      %37924 = func.call @_where_1050(%30619#1, %37923#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %37925 = arith.select %ptrue, %37923#0, %37923#0 : tensor<1x1x2048x32x64xbf16>
      %37926:2 = func.call @_einsum_1051(%31868, %31825, %37925) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4877)
      %37927 = arith.select %ptrue, %37926#1, %37926#1 : tensor<1x1x2048x32x64xbf16>
      %37928 = stablehlo.broadcast_in_dim %37927, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4878)
      %37929 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37930 = stablehlo.pad %37928, %37929, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4879)
      %37931 = stablehlo.convert %37926#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4880)
      %37932 = stablehlo.broadcast_in_dim %31867, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2541)
      %37933 = stablehlo.multiply %37931, %37932 : tensor<1x1x32x2048x2048xf32> loc(#loc2541)
      %37934 = stablehlo.multiply %37933, %31858 : tensor<1x1x32x2048x2048xf32> loc(#loc2541)
      %37935 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2539)
      %37936 = stablehlo.reduce(%37934 init: %37935) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2539)
        stablehlo.return %39096 : tensor<f32> loc(#loc2539)
      } loc(#loc2539)
      %37937 = stablehlo.reshape %37936 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4881)
      %37938 = stablehlo.negate %37937 : tensor<1x1x32x2048x1xf32> loc(#loc4882)
      %37939 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4883)
      %37940 = stablehlo.reduce(%37938 init: %37939) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4883)
        stablehlo.return %39096 : tensor<f32> loc(#loc4883)
      } loc(#loc4883)
      %37941 = stablehlo.broadcast_in_dim %37940, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4884)
      %37942 = stablehlo.broadcast_in_dim %37941, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4885)
      %37943 = stablehlo.broadcast_in_dim %31861, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2540)
      %37944 = stablehlo.divide %37931, %37943 : tensor<1x1x32x2048x2048xf32> loc(#loc2540)
      %37945 = stablehlo.add %37942, %37944 : tensor<1x1x32x2048x2048xf32> loc(#loc4886)
      %37946 = stablehlo.multiply %37945, %31858 : tensor<1x1x32x2048x2048xf32> loc(#loc2541)
      %37947 = func.call @_where_1052(%31852#1, %37946) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4887)
      %37948 = stablehlo.convert %37947 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4888)
      %37949 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37950 = stablehlo.broadcast_in_dim %37949, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2529)
      %37951 = stablehlo.multiply %37950, %37948 : tensor<1x1x32x2048x2048xbf16> loc(#loc2529)
      %37952 = stablehlo.multiply %37951, %31840 : tensor<1x1x32x2048x2048xbf16> loc(#loc2529)
      %37953 = stablehlo.multiply %37952, %31837 : tensor<1x1x32x2048x2048xbf16> loc(#loc2529)
      %37954 = stablehlo.add %37952, %37953 : tensor<1x1x32x2048x2048xbf16> loc(#loc4889)
      %37955 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %37956 = stablehlo.broadcast_in_dim %37955, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2526)
      %37957 = stablehlo.divide %37954, %37956 : tensor<1x1x32x2048x2048xbf16> loc(#loc2526)
      %37958:2 = func.call @_einsum_1053(%31824, %31832, %37957) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4890)
      %37959 = arith.select %ptrue, %37958#1, %37958#1 : tensor<1x1x2048x32x64xbf16>
      %37960 = stablehlo.broadcast_in_dim %37959, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4878)
      %37961 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37962 = stablehlo.pad %37960, %37961, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4891)
      %37963 = stablehlo.add %37930, %37962 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4892)
      %37964 = stablehlo.multiply %31823, %37958#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2524)
      %37965 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4893)
      %37966 = stablehlo.reduce(%37964 init: %37965) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4893)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4893)
      } loc(#loc4893)
      %37967 = stablehlo.reshape %37966 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4894)
      %37968 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4895)
      %37969 = stablehlo.reduce(%37967 init: %37968) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4895)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4895)
      } loc(#loc4895)
      %37970 = stablehlo.broadcast_in_dim %37969, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4896)
      %37971 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %37972 = stablehlo.broadcast_in_dim %37971, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2522)
      %37973 = stablehlo.multiply %37972, %37970 : tensor<1x64xbf16> loc(#loc2522)
      %37974 = func.call @softplus_1054(%31826#1, %37973) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4897)
      %37975 = func.call @_where_1055(%30618#1, %37974) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %37976 = stablehlo.broadcast_in_dim %31830, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2524)
      %37977 = stablehlo.multiply %37958#0, %37976 : tensor<1x1x2048x32x64xbf16> loc(#loc2524)
      %37978 = arith.select %ptrue, %37977, %37977 : tensor<1x1x2048x32x64xbf16>
      %37979 = stablehlo.broadcast_in_dim %37978, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4878)
      %37980 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %37981 = stablehlo.pad %37979, %37980, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4898)
      %37982 = stablehlo.add %37963, %37981 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4892)
      %37983:2 = func.call @_einsum_1056(%31815, %30617#0, %37982) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4899)
      %37984 = func.call @_where_1057(%30617#1, %37983#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %37985 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4900)
      %37986 = stablehlo.reduce(%37983#0 init: %37985) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4900)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4900)
      } loc(#loc4900)
      %37987 = stablehlo.reshape %37986 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4901)
      %37988 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4902)
      %37989 = stablehlo.reduce(%37987 init: %37988) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4902)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4902)
      } loc(#loc4902)
      %37990 = stablehlo.broadcast_in_dim %37989, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4903)
      %37991 = func.call @_where_1058(%30615#1, %37990) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %37992 = stablehlo.multiply %31806, %37983#0 : tensor<1x1x2048x2048xbf16> loc(#loc2512)
      %37993 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4904)
      %37994 = stablehlo.reduce(%37992 init: %37993) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4904)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4904)
      } loc(#loc4904)
      %37995 = stablehlo.reshape %37994 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4905)
      %37996 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4906)
      %37997 = stablehlo.reduce(%37995 init: %37996) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4906)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4906)
      } loc(#loc4906)
      %37998 = stablehlo.broadcast_in_dim %37997, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4907)
      %37999 = func.call @_where_1059(%30616#1, %37998) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38000 = stablehlo.broadcast_in_dim %31810, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2512)
      %38001 = stablehlo.multiply %37983#0, %38000 : tensor<1x1x2048x2048xbf16> loc(#loc2512)
      %38002 = stablehlo.multiply %31796, %38001 : tensor<1x1x2048x2048xbf16> loc(#loc2509)
      %38003 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4908)
      %38004 = stablehlo.reduce(%38002 init: %38003) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4908)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4908)
      } loc(#loc4908)
      %38005 = stablehlo.reshape %38004 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4909)
      %38006 = stablehlo.multiply %38005, %31804 : tensor<1x1x2048x1xbf16> loc(#loc2509)
      %38007 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38008 = stablehlo.broadcast_in_dim %38007, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2503)
      %38009 = stablehlo.divide %38006, %38008 : tensor<1x1x2048x1xbf16> loc(#loc2503)
      %38010 = stablehlo.convert %38009 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2499)
      %38011 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4910)
      %38012 = stablehlo.reduce(%38010 init: %38011) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4910)
        stablehlo.return %39096 : tensor<f32> loc(#loc4910)
      } loc(#loc4910)
      %38013 = stablehlo.broadcast_in_dim %38012, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4911)
      %38014 = stablehlo.broadcast_in_dim %38013, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4912)
      %38015 = stablehlo.convert %38014 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2502)
      %38016 = stablehlo.multiply %38015, %31786 : tensor<1x1x2048x2048xbf16> loc(#loc2505)
      %38017 = stablehlo.negate %38016 : tensor<1x1x2048x2048xbf16> loc(#loc4913)
      %38018 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4914)
      %38019 = stablehlo.reduce(%38017 init: %38018) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4914)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4914)
      } loc(#loc4914)
      %38020 = stablehlo.reshape %38019 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4915)
      %38021 = stablehlo.add %37921, %38016 : tensor<1x1x2048x2048xbf16> loc(#loc4916)
      %38022 = stablehlo.broadcast_in_dim %31800, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2509)
      %38023 = stablehlo.multiply %38001, %38022 : tensor<1x1x2048x2048xbf16> loc(#loc2509)
      %38024 = stablehlo.negate %38023 : tensor<1x1x2048x2048xbf16> loc(#loc4913)
      %38025 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4914)
      %38026 = stablehlo.reduce(%38024 init: %38025) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4914)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4914)
      } loc(#loc4914)
      %38027 = stablehlo.reshape %38026 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4915)
      %38028 = stablehlo.add %38021, %38023 : tensor<1x1x2048x2048xbf16> loc(#loc4916)
      %38029 = stablehlo.add %38020, %38027 : tensor<1x1x2048x1xbf16> loc(#loc4916)
      %38030 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38031 = stablehlo.broadcast_in_dim %38030, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2503)
      %38032 = stablehlo.divide %38029, %38031 : tensor<1x1x2048x1xbf16> loc(#loc2503)
      %38033 = stablehlo.convert %38032 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2499)
      %38034 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4910)
      %38035 = stablehlo.reduce(%38033 init: %38034) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4910)
        stablehlo.return %39096 : tensor<f32> loc(#loc4910)
      } loc(#loc4910)
      %38036 = stablehlo.broadcast_in_dim %38035, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4911)
      %38037 = stablehlo.broadcast_in_dim %38036, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4912)
      %38038 = stablehlo.convert %38037 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2502)
      %38039 = stablehlo.add %38028, %38038 : tensor<1x1x2048x2048xbf16> loc(#loc4917)
      %38040 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %38041 = stablehlo.broadcast_in_dim %38040, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2474)
      %38042 = stablehlo.multiply %38039, %38041 : tensor<1x1x2048x2048xbf16> loc(#loc2474)
      %38043 = stablehlo.broadcast_in_dim %31708, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2473)
      %38044 = stablehlo.multiply %38042, %38043 : tensor<1x1x2048x2048xbf16> loc(#loc2473)
      %38045 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4918)
      %38046 = stablehlo.reduce(%38044 init: %38045) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4918)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4918)
      } loc(#loc4918)
      %38047 = stablehlo.reshape %38046 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4919)
      %38048 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4920)
      %38049 = stablehlo.reduce(%38047 init: %38048) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4920)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4920)
      } loc(#loc4920)
      %38050 = stablehlo.broadcast_in_dim %38049, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4921)
      %38051 = func.call @_where_1060(%30600#1, %38050) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38052 = arith.select %ptrue, %38044, %38044 : tensor<1x1x2048x2048xbf16>
      %38053:2 = func.call @_einsum_1061(%30601#0, %31698, %38052) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4922)
      %38054 = func.call @_where_1062(%30601#1, %38053#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %38055 = stablehlo.broadcast_in_dim %31696, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2467)
      %38056 = stablehlo.multiply %38053#0, %38055 : tensor<1x1x2048x8192xbf16> loc(#loc2467)
      %38057 = stablehlo.multiply %31667, %38056 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38058 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %38059 = stablehlo.broadcast_in_dim %38058, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38060 = stablehlo.multiply %38059, %38057 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38061 = stablehlo.multiply %38060, %31684 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38062 = stablehlo.multiply %38061, %31681 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38063 = stablehlo.add %38061, %38062 : tensor<1x1x2048x8192xbf16> loc(#loc4923)
      %38064 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %38065 = stablehlo.broadcast_in_dim %38064, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38066 = stablehlo.multiply %38065, %38063 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38067 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %38068 = stablehlo.broadcast_in_dim %38067, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38069 = stablehlo.multiply %38068, %38066 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38070 = stablehlo.multiply %38069, %31673 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38071 = stablehlo.add %38066, %38070 : tensor<1x1x2048x8192xbf16> loc(#loc4923)
      %38072 = stablehlo.multiply %38056, %31690 : tensor<1x1x2048x8192xbf16> loc(#loc2460)
      %38073 = stablehlo.add %38071, %38072 : tensor<1x1x2048x8192xbf16> loc(#loc4923)
      %38074 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4924)
      %38075 = stablehlo.reduce(%38073 init: %38074) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4924)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4924)
      } loc(#loc4924)
      %38076 = stablehlo.reshape %38075 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4925)
      %38077 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4926)
      %38078 = stablehlo.reduce(%38076 init: %38077) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4926)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4926)
      } loc(#loc4926)
      %38079 = stablehlo.broadcast_in_dim %38078, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4927)
      %38080 = func.call @_where_1063(%30598#1, %38079) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %38081 = arith.select %ptrue, %38073, %38073 : tensor<1x1x2048x8192xbf16>
      %38082:2 = func.call @_einsum_1064(%30599#0, %31662, %38081) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc4928)
      %38083 = func.call @_where_1065(%30599#1, %38082#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %38084 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4929)
      %38085 = stablehlo.reduce(%38082#0 init: %38084) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4929)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4929)
      } loc(#loc4929)
      %38086 = stablehlo.reshape %38085 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4930)
      %38087 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4931)
      %38088 = stablehlo.reduce(%38086 init: %38087) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4931)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4931)
      } loc(#loc4931)
      %38089 = stablehlo.broadcast_in_dim %38088, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4932)
      %38090 = func.call @_where_1066(%30602#1, %38089) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38091 = stablehlo.multiply %31653, %38082#0 : tensor<1x1x2048x2048xbf16> loc(#loc2453)
      %38092 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4933)
      %38093 = stablehlo.reduce(%38091 init: %38092) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4933)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4933)
      } loc(#loc4933)
      %38094 = stablehlo.reshape %38093 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4934)
      %38095 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4935)
      %38096 = stablehlo.reduce(%38094 init: %38095) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4935)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4935)
      } loc(#loc4935)
      %38097 = stablehlo.broadcast_in_dim %38096, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4936)
      %38098 = func.call @_where_1067(%30603#1, %38097) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38099 = stablehlo.broadcast_in_dim %31657, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2453)
      %38100 = stablehlo.multiply %38082#0, %38099 : tensor<1x1x2048x2048xbf16> loc(#loc2453)
      %38101 = stablehlo.multiply %31643, %38100 : tensor<1x1x2048x2048xbf16> loc(#loc2450)
      %38102 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4937)
      %38103 = stablehlo.reduce(%38101 init: %38102) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4937)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4937)
      } loc(#loc4937)
      %38104 = stablehlo.reshape %38103 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4938)
      %38105 = stablehlo.multiply %38104, %31651 : tensor<1x1x2048x1xbf16> loc(#loc2450)
      %38106 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38107 = stablehlo.broadcast_in_dim %38106, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2444)
      %38108 = stablehlo.divide %38105, %38107 : tensor<1x1x2048x1xbf16> loc(#loc2444)
      %38109 = stablehlo.convert %38108 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2440)
      %38110 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4939)
      %38111 = stablehlo.reduce(%38109 init: %38110) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4939)
        stablehlo.return %39096 : tensor<f32> loc(#loc4939)
      } loc(#loc4939)
      %38112 = stablehlo.broadcast_in_dim %38111, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4940)
      %38113 = stablehlo.broadcast_in_dim %38112, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4941)
      %38114 = stablehlo.convert %38113 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2443)
      %38115 = stablehlo.multiply %38114, %31633 : tensor<1x1x2048x2048xbf16> loc(#loc2446)
      %38116 = stablehlo.negate %38115 : tensor<1x1x2048x2048xbf16> loc(#loc4942)
      %38117 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4943)
      %38118 = stablehlo.reduce(%38116 init: %38117) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4943)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4943)
      } loc(#loc4943)
      %38119 = stablehlo.reshape %38118 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4944)
      %38120 = stablehlo.add %38039, %38115 : tensor<1x1x2048x2048xbf16> loc(#loc4945)
      %38121 = stablehlo.broadcast_in_dim %31647, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2450)
      %38122 = stablehlo.multiply %38100, %38121 : tensor<1x1x2048x2048xbf16> loc(#loc2450)
      %38123 = stablehlo.negate %38122 : tensor<1x1x2048x2048xbf16> loc(#loc4942)
      %38124 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4943)
      %38125 = stablehlo.reduce(%38123 init: %38124) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4943)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4943)
      } loc(#loc4943)
      %38126 = stablehlo.reshape %38125 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4944)
      %38127 = stablehlo.add %38120, %38122 : tensor<1x1x2048x2048xbf16> loc(#loc4945)
      %38128 = stablehlo.add %38119, %38126 : tensor<1x1x2048x1xbf16> loc(#loc4945)
      %38129 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38130 = stablehlo.broadcast_in_dim %38129, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2444)
      %38131 = stablehlo.divide %38128, %38130 : tensor<1x1x2048x1xbf16> loc(#loc2444)
      %38132 = stablehlo.convert %38131 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2440)
      %38133 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4939)
      %38134 = stablehlo.reduce(%38132 init: %38133) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4939)
        stablehlo.return %39096 : tensor<f32> loc(#loc4939)
      } loc(#loc4939)
      %38135 = stablehlo.broadcast_in_dim %38134, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4940)
      %38136 = stablehlo.broadcast_in_dim %38135, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4941)
      %38137 = stablehlo.convert %38136 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2443)
      %38138 = stablehlo.add %38127, %38137 : tensor<1x1x2048x2048xbf16> loc(#loc4946)
      %38139 = stablehlo.add %35050, %38138 : tensor<1x1x2048x2048xbf16> loc(#loc4947)
      %38140 = arith.select %ptrue, %38138, %38138 : tensor<1x1x2048x2048xbf16>
      %38141:2 = func.call @_einsum_1068(%30608#0, %31615, %38140) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc4948)
      %38142 = func.call @_where_1069(%30608#1, %38141#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %38143 = arith.select %ptrue, %38141#0, %38141#0 : tensor<1x1x2048x32x64xbf16>
      %38144:2 = func.call @_einsum_1070(%31613, %31570, %38143) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4949)
      %38145 = arith.select %ptrue, %38144#1, %38144#1 : tensor<1x1x2048x32x64xbf16>
      %38146 = stablehlo.broadcast_in_dim %38145, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4950)
      %38147 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38148 = stablehlo.pad %38146, %38147, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4951)
      %38149 = stablehlo.convert %38144#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4952)
      %38150 = stablehlo.broadcast_in_dim %31612, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2431)
      %38151 = stablehlo.multiply %38149, %38150 : tensor<1x1x32x2048x2048xf32> loc(#loc2431)
      %38152 = stablehlo.multiply %38151, %31603 : tensor<1x1x32x2048x2048xf32> loc(#loc2431)
      %38153 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2429)
      %38154 = stablehlo.reduce(%38152 init: %38153) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2429)
        stablehlo.return %39096 : tensor<f32> loc(#loc2429)
      } loc(#loc2429)
      %38155 = stablehlo.reshape %38154 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc4953)
      %38156 = stablehlo.negate %38155 : tensor<1x1x32x2048x1xf32> loc(#loc4954)
      %38157 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4955)
      %38158 = stablehlo.reduce(%38156 init: %38157) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4955)
        stablehlo.return %39096 : tensor<f32> loc(#loc4955)
      } loc(#loc4955)
      %38159 = stablehlo.broadcast_in_dim %38158, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc4956)
      %38160 = stablehlo.broadcast_in_dim %38159, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4957)
      %38161 = stablehlo.broadcast_in_dim %31606, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2430)
      %38162 = stablehlo.divide %38149, %38161 : tensor<1x1x32x2048x2048xf32> loc(#loc2430)
      %38163 = stablehlo.add %38160, %38162 : tensor<1x1x32x2048x2048xf32> loc(#loc4958)
      %38164 = stablehlo.multiply %38163, %31603 : tensor<1x1x32x2048x2048xf32> loc(#loc2431)
      %38165 = func.call @_where_1071(%31597#1, %38164) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc4959)
      %38166 = stablehlo.convert %38165 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc4960)
      %38167 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38168 = stablehlo.broadcast_in_dim %38167, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2419)
      %38169 = stablehlo.multiply %38168, %38166 : tensor<1x1x32x2048x2048xbf16> loc(#loc2419)
      %38170 = stablehlo.multiply %38169, %31585 : tensor<1x1x32x2048x2048xbf16> loc(#loc2419)
      %38171 = stablehlo.multiply %38170, %31582 : tensor<1x1x32x2048x2048xbf16> loc(#loc2419)
      %38172 = stablehlo.add %38170, %38171 : tensor<1x1x32x2048x2048xbf16> loc(#loc4961)
      %38173 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38174 = stablehlo.broadcast_in_dim %38173, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2416)
      %38175 = stablehlo.divide %38172, %38174 : tensor<1x1x32x2048x2048xbf16> loc(#loc2416)
      %38176:2 = func.call @_einsum_1072(%31569, %31577, %38175) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc4962)
      %38177 = arith.select %ptrue, %38176#1, %38176#1 : tensor<1x1x2048x32x64xbf16>
      %38178 = stablehlo.broadcast_in_dim %38177, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4950)
      %38179 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38180 = stablehlo.pad %38178, %38179, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4963)
      %38181 = stablehlo.add %38148, %38180 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4964)
      %38182 = stablehlo.multiply %31568, %38176#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2414)
      %38183 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4965)
      %38184 = stablehlo.reduce(%38182 init: %38183) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4965)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4965)
      } loc(#loc4965)
      %38185 = stablehlo.reshape %38184 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc4966)
      %38186 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4967)
      %38187 = stablehlo.reduce(%38185 init: %38186) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4967)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4967)
      } loc(#loc4967)
      %38188 = stablehlo.broadcast_in_dim %38187, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc4968)
      %38189 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %38190 = stablehlo.broadcast_in_dim %38189, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2412)
      %38191 = stablehlo.multiply %38190, %38188 : tensor<1x64xbf16> loc(#loc2412)
      %38192 = func.call @softplus_1073(%31571#1, %38191) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4969)
      %38193 = func.call @_where_1074(%30607#1, %38192) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %38194 = stablehlo.broadcast_in_dim %31575, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2414)
      %38195 = stablehlo.multiply %38176#0, %38194 : tensor<1x1x2048x32x64xbf16> loc(#loc2414)
      %38196 = arith.select %ptrue, %38195, %38195 : tensor<1x1x2048x32x64xbf16>
      %38197 = stablehlo.broadcast_in_dim %38196, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc4950)
      %38198 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38199 = stablehlo.pad %38197, %38198, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc4970)
      %38200 = stablehlo.add %38181, %38199 : tensor<1x3x1x2048x32x64xbf16> loc(#loc4964)
      %38201:2 = func.call @_einsum_1075(%31560, %30606#0, %38200) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc4971)
      %38202 = func.call @_where_1076(%30606#1, %38201#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %38203 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4972)
      %38204 = stablehlo.reduce(%38201#0 init: %38203) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4972)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4972)
      } loc(#loc4972)
      %38205 = stablehlo.reshape %38204 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4973)
      %38206 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4974)
      %38207 = stablehlo.reduce(%38205 init: %38206) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4974)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4974)
      } loc(#loc4974)
      %38208 = stablehlo.broadcast_in_dim %38207, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4975)
      %38209 = func.call @_where_1077(%30604#1, %38208) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38210 = stablehlo.multiply %31551, %38201#0 : tensor<1x1x2048x2048xbf16> loc(#loc2402)
      %38211 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4976)
      %38212 = stablehlo.reduce(%38210 init: %38211) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4976)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4976)
      } loc(#loc4976)
      %38213 = stablehlo.reshape %38212 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4977)
      %38214 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4978)
      %38215 = stablehlo.reduce(%38213 init: %38214) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4978)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4978)
      } loc(#loc4978)
      %38216 = stablehlo.broadcast_in_dim %38215, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4979)
      %38217 = func.call @_where_1078(%30605#1, %38216) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38218 = stablehlo.broadcast_in_dim %31555, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2402)
      %38219 = stablehlo.multiply %38201#0, %38218 : tensor<1x1x2048x2048xbf16> loc(#loc2402)
      %38220 = stablehlo.multiply %31541, %38219 : tensor<1x1x2048x2048xbf16> loc(#loc2399)
      %38221 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4980)
      %38222 = stablehlo.reduce(%38220 init: %38221) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4980)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4980)
      } loc(#loc4980)
      %38223 = stablehlo.reshape %38222 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4981)
      %38224 = stablehlo.multiply %38223, %31549 : tensor<1x1x2048x1xbf16> loc(#loc2399)
      %38225 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38226 = stablehlo.broadcast_in_dim %38225, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2393)
      %38227 = stablehlo.divide %38224, %38226 : tensor<1x1x2048x1xbf16> loc(#loc2393)
      %38228 = stablehlo.convert %38227 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2389)
      %38229 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4982)
      %38230 = stablehlo.reduce(%38228 init: %38229) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4982)
        stablehlo.return %39096 : tensor<f32> loc(#loc4982)
      } loc(#loc4982)
      %38231 = stablehlo.broadcast_in_dim %38230, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4983)
      %38232 = stablehlo.broadcast_in_dim %38231, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4984)
      %38233 = stablehlo.convert %38232 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2392)
      %38234 = stablehlo.multiply %38233, %31531 : tensor<1x1x2048x2048xbf16> loc(#loc2395)
      %38235 = stablehlo.negate %38234 : tensor<1x1x2048x2048xbf16> loc(#loc4985)
      %38236 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4986)
      %38237 = stablehlo.reduce(%38235 init: %38236) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4986)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4986)
      } loc(#loc4986)
      %38238 = stablehlo.reshape %38237 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4987)
      %38239 = stablehlo.add %38139, %38234 : tensor<1x1x2048x2048xbf16> loc(#loc4988)
      %38240 = stablehlo.broadcast_in_dim %31545, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2399)
      %38241 = stablehlo.multiply %38219, %38240 : tensor<1x1x2048x2048xbf16> loc(#loc2399)
      %38242 = stablehlo.negate %38241 : tensor<1x1x2048x2048xbf16> loc(#loc4985)
      %38243 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4986)
      %38244 = stablehlo.reduce(%38242 init: %38243) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4986)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4986)
      } loc(#loc4986)
      %38245 = stablehlo.reshape %38244 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc4987)
      %38246 = stablehlo.add %38239, %38241 : tensor<1x1x2048x2048xbf16> loc(#loc4988)
      %38247 = stablehlo.add %38238, %38245 : tensor<1x1x2048x1xbf16> loc(#loc4988)
      %38248 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38249 = stablehlo.broadcast_in_dim %38248, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2393)
      %38250 = stablehlo.divide %38247, %38249 : tensor<1x1x2048x1xbf16> loc(#loc2393)
      %38251 = stablehlo.convert %38250 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2389)
      %38252 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4982)
      %38253 = stablehlo.reduce(%38251 init: %38252) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc4982)
        stablehlo.return %39096 : tensor<f32> loc(#loc4982)
      } loc(#loc4982)
      %38254 = stablehlo.broadcast_in_dim %38253, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc4983)
      %38255 = stablehlo.broadcast_in_dim %38254, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc4984)
      %38256 = stablehlo.convert %38255 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2392)
      %38257 = stablehlo.add %38246, %38256 : tensor<1x1x2048x2048xbf16> loc(#loc4989)
      %38258 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %38259 = stablehlo.broadcast_in_dim %38258, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2364)
      %38260 = stablehlo.multiply %38257, %38259 : tensor<1x1x2048x2048xbf16> loc(#loc2364)
      %38261 = stablehlo.broadcast_in_dim %31453, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2363)
      %38262 = stablehlo.multiply %38260, %38261 : tensor<1x1x2048x2048xbf16> loc(#loc2363)
      %38263 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4990)
      %38264 = stablehlo.reduce(%38262 init: %38263) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4990)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4990)
      } loc(#loc4990)
      %38265 = stablehlo.reshape %38264 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc4991)
      %38266 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4992)
      %38267 = stablehlo.reduce(%38265 init: %38266) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4992)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4992)
      } loc(#loc4992)
      %38268 = stablehlo.broadcast_in_dim %38267, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4993)
      %38269 = func.call @_where_1079(%30589#1, %38268) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38270 = arith.select %ptrue, %38262, %38262 : tensor<1x1x2048x2048xbf16>
      %38271:2 = func.call @_einsum_1080(%30590#0, %31443, %38270) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc4994)
      %38272 = func.call @_where_1081(%30590#1, %38271#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %38273 = stablehlo.broadcast_in_dim %31441, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2357)
      %38274 = stablehlo.multiply %38271#0, %38273 : tensor<1x1x2048x8192xbf16> loc(#loc2357)
      %38275 = stablehlo.multiply %31412, %38274 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38276 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %38277 = stablehlo.broadcast_in_dim %38276, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38278 = stablehlo.multiply %38277, %38275 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38279 = stablehlo.multiply %38278, %31429 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38280 = stablehlo.multiply %38279, %31426 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38281 = stablehlo.add %38279, %38280 : tensor<1x1x2048x8192xbf16> loc(#loc4995)
      %38282 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %38283 = stablehlo.broadcast_in_dim %38282, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38284 = stablehlo.multiply %38283, %38281 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38285 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %38286 = stablehlo.broadcast_in_dim %38285, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38287 = stablehlo.multiply %38286, %38284 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38288 = stablehlo.multiply %38287, %31418 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38289 = stablehlo.add %38284, %38288 : tensor<1x1x2048x8192xbf16> loc(#loc4995)
      %38290 = stablehlo.multiply %38274, %31435 : tensor<1x1x2048x8192xbf16> loc(#loc2350)
      %38291 = stablehlo.add %38289, %38290 : tensor<1x1x2048x8192xbf16> loc(#loc4995)
      %38292 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4996)
      %38293 = stablehlo.reduce(%38291 init: %38292) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4996)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4996)
      } loc(#loc4996)
      %38294 = stablehlo.reshape %38293 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc4997)
      %38295 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4998)
      %38296 = stablehlo.reduce(%38294 init: %38295) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc4998)
        stablehlo.return %39096 : tensor<bf16> loc(#loc4998)
      } loc(#loc4998)
      %38297 = stablehlo.broadcast_in_dim %38296, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4999)
      %38298 = func.call @_where_1082(%30587#1, %38297) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %38299 = arith.select %ptrue, %38291, %38291 : tensor<1x1x2048x8192xbf16>
      %38300:2 = func.call @_einsum_1083(%30588#0, %31407, %38299) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc5000)
      %38301 = func.call @_where_1084(%30588#1, %38300#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %38302 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5001)
      %38303 = stablehlo.reduce(%38300#0 init: %38302) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5001)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5001)
      } loc(#loc5001)
      %38304 = stablehlo.reshape %38303 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5002)
      %38305 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5003)
      %38306 = stablehlo.reduce(%38304 init: %38305) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5003)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5003)
      } loc(#loc5003)
      %38307 = stablehlo.broadcast_in_dim %38306, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5004)
      %38308 = func.call @_where_1085(%30591#1, %38307) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38309 = stablehlo.multiply %31398, %38300#0 : tensor<1x1x2048x2048xbf16> loc(#loc2343)
      %38310 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5005)
      %38311 = stablehlo.reduce(%38309 init: %38310) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5005)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5005)
      } loc(#loc5005)
      %38312 = stablehlo.reshape %38311 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5006)
      %38313 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5007)
      %38314 = stablehlo.reduce(%38312 init: %38313) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5007)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5007)
      } loc(#loc5007)
      %38315 = stablehlo.broadcast_in_dim %38314, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5008)
      %38316 = func.call @_where_1086(%30592#1, %38315) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38317 = stablehlo.broadcast_in_dim %31402, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2343)
      %38318 = stablehlo.multiply %38300#0, %38317 : tensor<1x1x2048x2048xbf16> loc(#loc2343)
      %38319 = stablehlo.multiply %31388, %38318 : tensor<1x1x2048x2048xbf16> loc(#loc2340)
      %38320 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5009)
      %38321 = stablehlo.reduce(%38319 init: %38320) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5009)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5009)
      } loc(#loc5009)
      %38322 = stablehlo.reshape %38321 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5010)
      %38323 = stablehlo.multiply %38322, %31396 : tensor<1x1x2048x1xbf16> loc(#loc2340)
      %38324 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38325 = stablehlo.broadcast_in_dim %38324, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2334)
      %38326 = stablehlo.divide %38323, %38325 : tensor<1x1x2048x1xbf16> loc(#loc2334)
      %38327 = stablehlo.convert %38326 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2330)
      %38328 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5011)
      %38329 = stablehlo.reduce(%38327 init: %38328) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5011)
        stablehlo.return %39096 : tensor<f32> loc(#loc5011)
      } loc(#loc5011)
      %38330 = stablehlo.broadcast_in_dim %38329, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5012)
      %38331 = stablehlo.broadcast_in_dim %38330, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5013)
      %38332 = stablehlo.convert %38331 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2333)
      %38333 = stablehlo.multiply %38332, %31378 : tensor<1x1x2048x2048xbf16> loc(#loc2336)
      %38334 = stablehlo.negate %38333 : tensor<1x1x2048x2048xbf16> loc(#loc5014)
      %38335 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5015)
      %38336 = stablehlo.reduce(%38334 init: %38335) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5015)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5015)
      } loc(#loc5015)
      %38337 = stablehlo.reshape %38336 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5016)
      %38338 = stablehlo.add %38257, %38333 : tensor<1x1x2048x2048xbf16> loc(#loc5017)
      %38339 = stablehlo.broadcast_in_dim %31392, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2340)
      %38340 = stablehlo.multiply %38318, %38339 : tensor<1x1x2048x2048xbf16> loc(#loc2340)
      %38341 = stablehlo.negate %38340 : tensor<1x1x2048x2048xbf16> loc(#loc5014)
      %38342 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5015)
      %38343 = stablehlo.reduce(%38341 init: %38342) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5015)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5015)
      } loc(#loc5015)
      %38344 = stablehlo.reshape %38343 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5016)
      %38345 = stablehlo.add %38338, %38340 : tensor<1x1x2048x2048xbf16> loc(#loc5017)
      %38346 = stablehlo.add %38337, %38344 : tensor<1x1x2048x1xbf16> loc(#loc5017)
      %38347 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38348 = stablehlo.broadcast_in_dim %38347, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2334)
      %38349 = stablehlo.divide %38346, %38348 : tensor<1x1x2048x1xbf16> loc(#loc2334)
      %38350 = stablehlo.convert %38349 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2330)
      %38351 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5011)
      %38352 = stablehlo.reduce(%38350 init: %38351) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5011)
        stablehlo.return %39096 : tensor<f32> loc(#loc5011)
      } loc(#loc5011)
      %38353 = stablehlo.broadcast_in_dim %38352, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5012)
      %38354 = stablehlo.broadcast_in_dim %38353, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5013)
      %38355 = stablehlo.convert %38354 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2333)
      %38356 = stablehlo.add %38345, %38355 : tensor<1x1x2048x2048xbf16> loc(#loc5018)
      %38357 = stablehlo.add %35091, %38356 : tensor<1x1x2048x2048xbf16> loc(#loc5019)
      %38358 = arith.select %ptrue, %38356, %38356 : tensor<1x1x2048x2048xbf16>
      %38359:2 = func.call @_einsum_1087(%30597#0, %31360, %38358) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc5020)
      %38360 = func.call @_where_1088(%30597#1, %38359#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %38361 = arith.select %ptrue, %38359#0, %38359#0 : tensor<1x1x2048x32x64xbf16>
      %38362:2 = func.call @_einsum_1089(%31358, %31315, %38361) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc5021)
      %38363 = arith.select %ptrue, %38362#1, %38362#1 : tensor<1x1x2048x32x64xbf16>
      %38364 = stablehlo.broadcast_in_dim %38363, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5022)
      %38365 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38366 = stablehlo.pad %38364, %38365, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5023)
      %38367 = stablehlo.convert %38362#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5024)
      %38368 = stablehlo.broadcast_in_dim %31357, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2321)
      %38369 = stablehlo.multiply %38367, %38368 : tensor<1x1x32x2048x2048xf32> loc(#loc2321)
      %38370 = stablehlo.multiply %38369, %31348 : tensor<1x1x32x2048x2048xf32> loc(#loc2321)
      %38371 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2319)
      %38372 = stablehlo.reduce(%38370 init: %38371) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2319)
        stablehlo.return %39096 : tensor<f32> loc(#loc2319)
      } loc(#loc2319)
      %38373 = stablehlo.reshape %38372 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc5025)
      %38374 = stablehlo.negate %38373 : tensor<1x1x32x2048x1xf32> loc(#loc5026)
      %38375 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5027)
      %38376 = stablehlo.reduce(%38374 init: %38375) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5027)
        stablehlo.return %39096 : tensor<f32> loc(#loc5027)
      } loc(#loc5027)
      %38377 = stablehlo.broadcast_in_dim %38376, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc5028)
      %38378 = stablehlo.broadcast_in_dim %38377, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5029)
      %38379 = stablehlo.broadcast_in_dim %31351, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2320)
      %38380 = stablehlo.divide %38367, %38379 : tensor<1x1x32x2048x2048xf32> loc(#loc2320)
      %38381 = stablehlo.add %38378, %38380 : tensor<1x1x32x2048x2048xf32> loc(#loc5030)
      %38382 = stablehlo.multiply %38381, %31348 : tensor<1x1x32x2048x2048xf32> loc(#loc2321)
      %38383 = func.call @_where_1090(%31342#1, %38382) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5031)
      %38384 = stablehlo.convert %38383 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5032)
      %38385 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38386 = stablehlo.broadcast_in_dim %38385, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2309)
      %38387 = stablehlo.multiply %38386, %38384 : tensor<1x1x32x2048x2048xbf16> loc(#loc2309)
      %38388 = stablehlo.multiply %38387, %31330 : tensor<1x1x32x2048x2048xbf16> loc(#loc2309)
      %38389 = stablehlo.multiply %38388, %31327 : tensor<1x1x32x2048x2048xbf16> loc(#loc2309)
      %38390 = stablehlo.add %38388, %38389 : tensor<1x1x32x2048x2048xbf16> loc(#loc5033)
      %38391 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38392 = stablehlo.broadcast_in_dim %38391, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2306)
      %38393 = stablehlo.divide %38390, %38392 : tensor<1x1x32x2048x2048xbf16> loc(#loc2306)
      %38394:2 = func.call @_einsum_1091(%31314, %31322, %38393) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc5034)
      %38395 = arith.select %ptrue, %38394#1, %38394#1 : tensor<1x1x2048x32x64xbf16>
      %38396 = stablehlo.broadcast_in_dim %38395, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5022)
      %38397 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38398 = stablehlo.pad %38396, %38397, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5035)
      %38399 = stablehlo.add %38366, %38398 : tensor<1x3x1x2048x32x64xbf16> loc(#loc5036)
      %38400 = stablehlo.multiply %31313, %38394#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2304)
      %38401 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5037)
      %38402 = stablehlo.reduce(%38400 init: %38401) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5037)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5037)
      } loc(#loc5037)
      %38403 = stablehlo.reshape %38402 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc5038)
      %38404 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5039)
      %38405 = stablehlo.reduce(%38403 init: %38404) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5039)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5039)
      } loc(#loc5039)
      %38406 = stablehlo.broadcast_in_dim %38405, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5040)
      %38407 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %38408 = stablehlo.broadcast_in_dim %38407, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2302)
      %38409 = stablehlo.multiply %38408, %38406 : tensor<1x64xbf16> loc(#loc2302)
      %38410 = func.call @softplus_1092(%31316#1, %38409) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc5041)
      %38411 = func.call @_where_1093(%30596#1, %38410) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %38412 = stablehlo.broadcast_in_dim %31320, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2304)
      %38413 = stablehlo.multiply %38394#0, %38412 : tensor<1x1x2048x32x64xbf16> loc(#loc2304)
      %38414 = arith.select %ptrue, %38413, %38413 : tensor<1x1x2048x32x64xbf16>
      %38415 = stablehlo.broadcast_in_dim %38414, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5022)
      %38416 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38417 = stablehlo.pad %38415, %38416, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5042)
      %38418 = stablehlo.add %38399, %38417 : tensor<1x3x1x2048x32x64xbf16> loc(#loc5036)
      %38419:2 = func.call @_einsum_1094(%31305, %30595#0, %38418) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc5043)
      %38420 = func.call @_where_1095(%30595#1, %38419#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %38421 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5044)
      %38422 = stablehlo.reduce(%38419#0 init: %38421) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5044)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5044)
      } loc(#loc5044)
      %38423 = stablehlo.reshape %38422 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5045)
      %38424 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5046)
      %38425 = stablehlo.reduce(%38423 init: %38424) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5046)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5046)
      } loc(#loc5046)
      %38426 = stablehlo.broadcast_in_dim %38425, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5047)
      %38427 = func.call @_where_1096(%30593#1, %38426) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38428 = stablehlo.multiply %31296, %38419#0 : tensor<1x1x2048x2048xbf16> loc(#loc2292)
      %38429 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5048)
      %38430 = stablehlo.reduce(%38428 init: %38429) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5048)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5048)
      } loc(#loc5048)
      %38431 = stablehlo.reshape %38430 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5049)
      %38432 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5050)
      %38433 = stablehlo.reduce(%38431 init: %38432) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5050)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5050)
      } loc(#loc5050)
      %38434 = stablehlo.broadcast_in_dim %38433, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5051)
      %38435 = func.call @_where_1097(%30594#1, %38434) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38436 = stablehlo.broadcast_in_dim %31300, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2292)
      %38437 = stablehlo.multiply %38419#0, %38436 : tensor<1x1x2048x2048xbf16> loc(#loc2292)
      %38438 = stablehlo.multiply %31286, %38437 : tensor<1x1x2048x2048xbf16> loc(#loc2289)
      %38439 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5052)
      %38440 = stablehlo.reduce(%38438 init: %38439) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5052)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5052)
      } loc(#loc5052)
      %38441 = stablehlo.reshape %38440 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5053)
      %38442 = stablehlo.multiply %38441, %31294 : tensor<1x1x2048x1xbf16> loc(#loc2289)
      %38443 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38444 = stablehlo.broadcast_in_dim %38443, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2283)
      %38445 = stablehlo.divide %38442, %38444 : tensor<1x1x2048x1xbf16> loc(#loc2283)
      %38446 = stablehlo.convert %38445 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2279)
      %38447 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5054)
      %38448 = stablehlo.reduce(%38446 init: %38447) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5054)
        stablehlo.return %39096 : tensor<f32> loc(#loc5054)
      } loc(#loc5054)
      %38449 = stablehlo.broadcast_in_dim %38448, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5055)
      %38450 = stablehlo.broadcast_in_dim %38449, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5056)
      %38451 = stablehlo.convert %38450 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2282)
      %38452 = stablehlo.multiply %38451, %31276 : tensor<1x1x2048x2048xbf16> loc(#loc2285)
      %38453 = stablehlo.negate %38452 : tensor<1x1x2048x2048xbf16> loc(#loc5057)
      %38454 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5058)
      %38455 = stablehlo.reduce(%38453 init: %38454) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5058)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5058)
      } loc(#loc5058)
      %38456 = stablehlo.reshape %38455 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5059)
      %38457 = stablehlo.add %38357, %38452 : tensor<1x1x2048x2048xbf16> loc(#loc5060)
      %38458 = stablehlo.broadcast_in_dim %31290, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2289)
      %38459 = stablehlo.multiply %38437, %38458 : tensor<1x1x2048x2048xbf16> loc(#loc2289)
      %38460 = stablehlo.negate %38459 : tensor<1x1x2048x2048xbf16> loc(#loc5057)
      %38461 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5058)
      %38462 = stablehlo.reduce(%38460 init: %38461) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5058)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5058)
      } loc(#loc5058)
      %38463 = stablehlo.reshape %38462 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5059)
      %38464 = stablehlo.add %38457, %38459 : tensor<1x1x2048x2048xbf16> loc(#loc5060)
      %38465 = stablehlo.add %38456, %38463 : tensor<1x1x2048x1xbf16> loc(#loc5060)
      %38466 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38467 = stablehlo.broadcast_in_dim %38466, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2283)
      %38468 = stablehlo.divide %38465, %38467 : tensor<1x1x2048x1xbf16> loc(#loc2283)
      %38469 = stablehlo.convert %38468 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2279)
      %38470 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5054)
      %38471 = stablehlo.reduce(%38469 init: %38470) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5054)
        stablehlo.return %39096 : tensor<f32> loc(#loc5054)
      } loc(#loc5054)
      %38472 = stablehlo.broadcast_in_dim %38471, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5055)
      %38473 = stablehlo.broadcast_in_dim %38472, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5056)
      %38474 = stablehlo.convert %38473 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2282)
      %38475 = stablehlo.add %38464, %38474 : tensor<1x1x2048x2048xbf16> loc(#loc5061)
      %38476 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %38477 = stablehlo.broadcast_in_dim %38476, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2254)
      %38478 = stablehlo.multiply %38475, %38477 : tensor<1x1x2048x2048xbf16> loc(#loc2254)
      %38479 = stablehlo.broadcast_in_dim %31198, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2253)
      %38480 = stablehlo.multiply %38478, %38479 : tensor<1x1x2048x2048xbf16> loc(#loc2253)
      %38481 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5062)
      %38482 = stablehlo.reduce(%38480 init: %38481) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5062)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5062)
      } loc(#loc5062)
      %38483 = stablehlo.reshape %38482 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5063)
      %38484 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5064)
      %38485 = stablehlo.reduce(%38483 init: %38484) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5064)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5064)
      } loc(#loc5064)
      %38486 = stablehlo.broadcast_in_dim %38485, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5065)
      %38487 = func.call @_where_1098(%30512#1, %38486) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38488 = arith.select %ptrue, %38480, %38480 : tensor<1x1x2048x2048xbf16>
      %38489:2 = func.call @_einsum_1099(%30513#0, %31188, %38488) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc5066)
      %38490 = func.call @_where_1100(%30513#1, %38489#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %38491 = stablehlo.broadcast_in_dim %31186, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2247)
      %38492 = stablehlo.multiply %38489#0, %38491 : tensor<1x1x2048x8192xbf16> loc(#loc2247)
      %38493 = stablehlo.multiply %31157, %38492 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38494 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %38495 = stablehlo.broadcast_in_dim %38494, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38496 = stablehlo.multiply %38495, %38493 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38497 = stablehlo.multiply %38496, %31174 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38498 = stablehlo.multiply %38497, %31171 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38499 = stablehlo.add %38497, %38498 : tensor<1x1x2048x8192xbf16> loc(#loc5067)
      %38500 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %38501 = stablehlo.broadcast_in_dim %38500, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38502 = stablehlo.multiply %38501, %38499 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38503 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %38504 = stablehlo.broadcast_in_dim %38503, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38505 = stablehlo.multiply %38504, %38502 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38506 = stablehlo.multiply %38505, %31163 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38507 = stablehlo.add %38502, %38506 : tensor<1x1x2048x8192xbf16> loc(#loc5067)
      %38508 = stablehlo.multiply %38492, %31180 : tensor<1x1x2048x8192xbf16> loc(#loc2240)
      %38509 = stablehlo.add %38507, %38508 : tensor<1x1x2048x8192xbf16> loc(#loc5067)
      %38510 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5068)
      %38511 = stablehlo.reduce(%38509 init: %38510) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5068)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5068)
      } loc(#loc5068)
      %38512 = stablehlo.reshape %38511 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc5069)
      %38513 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5070)
      %38514 = stablehlo.reduce(%38512 init: %38513) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5070)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5070)
      } loc(#loc5070)
      %38515 = stablehlo.broadcast_in_dim %38514, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc5071)
      %38516 = func.call @_where_1101(%30510#1, %38515) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %38517 = arith.select %ptrue, %38509, %38509 : tensor<1x1x2048x8192xbf16>
      %38518:2 = func.call @_einsum_1102(%30511#0, %31152, %38517) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc5072)
      %38519 = func.call @_where_1103(%30511#1, %38518#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %38520 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5073)
      %38521 = stablehlo.reduce(%38518#0 init: %38520) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5073)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5073)
      } loc(#loc5073)
      %38522 = stablehlo.reshape %38521 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5074)
      %38523 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5075)
      %38524 = stablehlo.reduce(%38522 init: %38523) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5075)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5075)
      } loc(#loc5075)
      %38525 = stablehlo.broadcast_in_dim %38524, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5076)
      %38526 = func.call @_where_1104(%30514#1, %38525) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38527 = stablehlo.multiply %31143, %38518#0 : tensor<1x1x2048x2048xbf16> loc(#loc2233)
      %38528 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5077)
      %38529 = stablehlo.reduce(%38527 init: %38528) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5077)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5077)
      } loc(#loc5077)
      %38530 = stablehlo.reshape %38529 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5078)
      %38531 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5079)
      %38532 = stablehlo.reduce(%38530 init: %38531) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5079)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5079)
      } loc(#loc5079)
      %38533 = stablehlo.broadcast_in_dim %38532, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5080)
      %38534 = func.call @_where_1105(%30515#1, %38533) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38535 = stablehlo.broadcast_in_dim %31147, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2233)
      %38536 = stablehlo.multiply %38518#0, %38535 : tensor<1x1x2048x2048xbf16> loc(#loc2233)
      %38537 = stablehlo.multiply %31133, %38536 : tensor<1x1x2048x2048xbf16> loc(#loc2230)
      %38538 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5081)
      %38539 = stablehlo.reduce(%38537 init: %38538) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5081)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5081)
      } loc(#loc5081)
      %38540 = stablehlo.reshape %38539 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5082)
      %38541 = stablehlo.multiply %38540, %31141 : tensor<1x1x2048x1xbf16> loc(#loc2230)
      %38542 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38543 = stablehlo.broadcast_in_dim %38542, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2224)
      %38544 = stablehlo.divide %38541, %38543 : tensor<1x1x2048x1xbf16> loc(#loc2224)
      %38545 = stablehlo.convert %38544 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2220)
      %38546 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5083)
      %38547 = stablehlo.reduce(%38545 init: %38546) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5083)
        stablehlo.return %39096 : tensor<f32> loc(#loc5083)
      } loc(#loc5083)
      %38548 = stablehlo.broadcast_in_dim %38547, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5084)
      %38549 = stablehlo.broadcast_in_dim %38548, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5085)
      %38550 = stablehlo.convert %38549 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2223)
      %38551 = stablehlo.multiply %38550, %31123 : tensor<1x1x2048x2048xbf16> loc(#loc2226)
      %38552 = stablehlo.negate %38551 : tensor<1x1x2048x2048xbf16> loc(#loc5086)
      %38553 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5087)
      %38554 = stablehlo.reduce(%38552 init: %38553) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5087)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5087)
      } loc(#loc5087)
      %38555 = stablehlo.reshape %38554 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5088)
      %38556 = stablehlo.add %38475, %38551 : tensor<1x1x2048x2048xbf16> loc(#loc5089)
      %38557 = stablehlo.broadcast_in_dim %31137, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2230)
      %38558 = stablehlo.multiply %38536, %38557 : tensor<1x1x2048x2048xbf16> loc(#loc2230)
      %38559 = stablehlo.negate %38558 : tensor<1x1x2048x2048xbf16> loc(#loc5086)
      %38560 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5087)
      %38561 = stablehlo.reduce(%38559 init: %38560) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5087)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5087)
      } loc(#loc5087)
      %38562 = stablehlo.reshape %38561 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5088)
      %38563 = stablehlo.add %38556, %38558 : tensor<1x1x2048x2048xbf16> loc(#loc5089)
      %38564 = stablehlo.add %38555, %38562 : tensor<1x1x2048x1xbf16> loc(#loc5089)
      %38565 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38566 = stablehlo.broadcast_in_dim %38565, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2224)
      %38567 = stablehlo.divide %38564, %38566 : tensor<1x1x2048x1xbf16> loc(#loc2224)
      %38568 = stablehlo.convert %38567 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2220)
      %38569 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5083)
      %38570 = stablehlo.reduce(%38568 init: %38569) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5083)
        stablehlo.return %39096 : tensor<f32> loc(#loc5083)
      } loc(#loc5083)
      %38571 = stablehlo.broadcast_in_dim %38570, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5084)
      %38572 = stablehlo.broadcast_in_dim %38571, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5085)
      %38573 = stablehlo.convert %38572 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2223)
      %38574 = stablehlo.add %38563, %38573 : tensor<1x1x2048x2048xbf16> loc(#loc5090)
      %38575 = stablehlo.add %35378, %38574 : tensor<1x1x2048x2048xbf16> loc(#loc5091)
      %38576 = arith.select %ptrue, %38574, %38574 : tensor<1x1x2048x2048xbf16>
      %38577:2 = func.call @_einsum_1106(%30520#0, %31105, %38576) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc5092)
      %38578 = func.call @_where_1107(%30520#1, %38577#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %38579 = arith.select %ptrue, %38577#0, %38577#0 : tensor<1x1x2048x32x64xbf16>
      %38580:2 = func.call @_einsum_1108(%31103, %31060, %38579) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc5093)
      %38581 = arith.select %ptrue, %38580#1, %38580#1 : tensor<1x1x2048x32x64xbf16>
      %38582 = stablehlo.broadcast_in_dim %38581, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5094)
      %38583 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38584 = stablehlo.pad %38582, %38583, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5095)
      %38585 = stablehlo.convert %38580#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5096)
      %38586 = stablehlo.broadcast_in_dim %31102, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2211)
      %38587 = stablehlo.multiply %38585, %38586 : tensor<1x1x32x2048x2048xf32> loc(#loc2211)
      %38588 = stablehlo.multiply %38587, %31093 : tensor<1x1x32x2048x2048xf32> loc(#loc2211)
      %38589 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2209)
      %38590 = stablehlo.reduce(%38588 init: %38589) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2209)
        stablehlo.return %39096 : tensor<f32> loc(#loc2209)
      } loc(#loc2209)
      %38591 = stablehlo.reshape %38590 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc5097)
      %38592 = stablehlo.negate %38591 : tensor<1x1x32x2048x1xf32> loc(#loc5098)
      %38593 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5099)
      %38594 = stablehlo.reduce(%38592 init: %38593) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5099)
        stablehlo.return %39096 : tensor<f32> loc(#loc5099)
      } loc(#loc5099)
      %38595 = stablehlo.broadcast_in_dim %38594, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc5100)
      %38596 = stablehlo.broadcast_in_dim %38595, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5101)
      %38597 = stablehlo.broadcast_in_dim %31096, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2210)
      %38598 = stablehlo.divide %38585, %38597 : tensor<1x1x32x2048x2048xf32> loc(#loc2210)
      %38599 = stablehlo.add %38596, %38598 : tensor<1x1x32x2048x2048xf32> loc(#loc5102)
      %38600 = stablehlo.multiply %38599, %31093 : tensor<1x1x32x2048x2048xf32> loc(#loc2211)
      %38601 = func.call @_where_1109(%31087#1, %38600) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5103)
      %38602 = stablehlo.convert %38601 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5104)
      %38603 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38604 = stablehlo.broadcast_in_dim %38603, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2199)
      %38605 = stablehlo.multiply %38604, %38602 : tensor<1x1x32x2048x2048xbf16> loc(#loc2199)
      %38606 = stablehlo.multiply %38605, %31075 : tensor<1x1x32x2048x2048xbf16> loc(#loc2199)
      %38607 = stablehlo.multiply %38606, %31072 : tensor<1x1x32x2048x2048xbf16> loc(#loc2199)
      %38608 = stablehlo.add %38606, %38607 : tensor<1x1x32x2048x2048xbf16> loc(#loc5105)
      %38609 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38610 = stablehlo.broadcast_in_dim %38609, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2196)
      %38611 = stablehlo.divide %38608, %38610 : tensor<1x1x32x2048x2048xbf16> loc(#loc2196)
      %38612:2 = func.call @_einsum_1110(%31059, %31067, %38611) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc5106)
      %38613 = arith.select %ptrue, %38612#1, %38612#1 : tensor<1x1x2048x32x64xbf16>
      %38614 = stablehlo.broadcast_in_dim %38613, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5094)
      %38615 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38616 = stablehlo.pad %38614, %38615, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5107)
      %38617 = stablehlo.add %38584, %38616 : tensor<1x3x1x2048x32x64xbf16> loc(#loc5108)
      %38618 = stablehlo.multiply %31058, %38612#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2194)
      %38619 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5109)
      %38620 = stablehlo.reduce(%38618 init: %38619) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5109)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5109)
      } loc(#loc5109)
      %38621 = stablehlo.reshape %38620 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc5110)
      %38622 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5111)
      %38623 = stablehlo.reduce(%38621 init: %38622) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5111)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5111)
      } loc(#loc5111)
      %38624 = stablehlo.broadcast_in_dim %38623, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5112)
      %38625 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %38626 = stablehlo.broadcast_in_dim %38625, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2192)
      %38627 = stablehlo.multiply %38626, %38624 : tensor<1x64xbf16> loc(#loc2192)
      %38628 = func.call @softplus_1111(%31061#1, %38627) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc5113)
      %38629 = func.call @_where_1112(%30519#1, %38628) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %38630 = stablehlo.broadcast_in_dim %31065, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2194)
      %38631 = stablehlo.multiply %38612#0, %38630 : tensor<1x1x2048x32x64xbf16> loc(#loc2194)
      %38632 = arith.select %ptrue, %38631, %38631 : tensor<1x1x2048x32x64xbf16>
      %38633 = stablehlo.broadcast_in_dim %38632, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5094)
      %38634 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38635 = stablehlo.pad %38633, %38634, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5114)
      %38636 = stablehlo.add %38617, %38635 : tensor<1x3x1x2048x32x64xbf16> loc(#loc5108)
      %38637:2 = func.call @_einsum_1113(%31050, %30518#0, %38636) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc5115)
      %38638 = func.call @_where_1114(%30518#1, %38637#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %38639 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5116)
      %38640 = stablehlo.reduce(%38637#0 init: %38639) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5116)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5116)
      } loc(#loc5116)
      %38641 = stablehlo.reshape %38640 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5117)
      %38642 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5118)
      %38643 = stablehlo.reduce(%38641 init: %38642) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5118)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5118)
      } loc(#loc5118)
      %38644 = stablehlo.broadcast_in_dim %38643, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5119)
      %38645 = func.call @_where_1115(%30516#1, %38644) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38646 = stablehlo.multiply %31041, %38637#0 : tensor<1x1x2048x2048xbf16> loc(#loc2182)
      %38647 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5120)
      %38648 = stablehlo.reduce(%38646 init: %38647) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5120)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5120)
      } loc(#loc5120)
      %38649 = stablehlo.reshape %38648 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5121)
      %38650 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5122)
      %38651 = stablehlo.reduce(%38649 init: %38650) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5122)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5122)
      } loc(#loc5122)
      %38652 = stablehlo.broadcast_in_dim %38651, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5123)
      %38653 = func.call @_where_1116(%30517#1, %38652) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38654 = stablehlo.broadcast_in_dim %31045, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2182)
      %38655 = stablehlo.multiply %38637#0, %38654 : tensor<1x1x2048x2048xbf16> loc(#loc2182)
      %38656 = stablehlo.multiply %31031, %38655 : tensor<1x1x2048x2048xbf16> loc(#loc2179)
      %38657 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5124)
      %38658 = stablehlo.reduce(%38656 init: %38657) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5124)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5124)
      } loc(#loc5124)
      %38659 = stablehlo.reshape %38658 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5125)
      %38660 = stablehlo.multiply %38659, %31039 : tensor<1x1x2048x1xbf16> loc(#loc2179)
      %38661 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38662 = stablehlo.broadcast_in_dim %38661, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2173)
      %38663 = stablehlo.divide %38660, %38662 : tensor<1x1x2048x1xbf16> loc(#loc2173)
      %38664 = stablehlo.convert %38663 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2169)
      %38665 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5126)
      %38666 = stablehlo.reduce(%38664 init: %38665) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5126)
        stablehlo.return %39096 : tensor<f32> loc(#loc5126)
      } loc(#loc5126)
      %38667 = stablehlo.broadcast_in_dim %38666, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5127)
      %38668 = stablehlo.broadcast_in_dim %38667, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5128)
      %38669 = stablehlo.convert %38668 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2172)
      %38670 = stablehlo.multiply %38669, %31021 : tensor<1x1x2048x2048xbf16> loc(#loc2175)
      %38671 = stablehlo.negate %38670 : tensor<1x1x2048x2048xbf16> loc(#loc5129)
      %38672 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5130)
      %38673 = stablehlo.reduce(%38671 init: %38672) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5130)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5130)
      } loc(#loc5130)
      %38674 = stablehlo.reshape %38673 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5131)
      %38675 = stablehlo.add %38575, %38670 : tensor<1x1x2048x2048xbf16> loc(#loc5132)
      %38676 = stablehlo.broadcast_in_dim %31035, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2179)
      %38677 = stablehlo.multiply %38655, %38676 : tensor<1x1x2048x2048xbf16> loc(#loc2179)
      %38678 = stablehlo.negate %38677 : tensor<1x1x2048x2048xbf16> loc(#loc5129)
      %38679 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5130)
      %38680 = stablehlo.reduce(%38678 init: %38679) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5130)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5130)
      } loc(#loc5130)
      %38681 = stablehlo.reshape %38680 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5131)
      %38682 = stablehlo.add %38675, %38677 : tensor<1x1x2048x2048xbf16> loc(#loc5132)
      %38683 = stablehlo.add %38674, %38681 : tensor<1x1x2048x1xbf16> loc(#loc5132)
      %38684 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38685 = stablehlo.broadcast_in_dim %38684, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2173)
      %38686 = stablehlo.divide %38683, %38685 : tensor<1x1x2048x1xbf16> loc(#loc2173)
      %38687 = stablehlo.convert %38686 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2169)
      %38688 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5126)
      %38689 = stablehlo.reduce(%38687 init: %38688) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5126)
        stablehlo.return %39096 : tensor<f32> loc(#loc5126)
      } loc(#loc5126)
      %38690 = stablehlo.broadcast_in_dim %38689, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5127)
      %38691 = stablehlo.broadcast_in_dim %38690, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5128)
      %38692 = stablehlo.convert %38691 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2172)
      %38693 = stablehlo.add %38682, %38692 : tensor<1x1x2048x2048xbf16> loc(#loc5133)
      %38694 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc1984)
      %38695 = stablehlo.broadcast_in_dim %38694, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2144)
      %38696 = stablehlo.multiply %38693, %38695 : tensor<1x1x2048x2048xbf16> loc(#loc2144)
      %38697 = stablehlo.broadcast_in_dim %30943, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2143)
      %38698 = stablehlo.multiply %38696, %38697 : tensor<1x1x2048x2048xbf16> loc(#loc2143)
      %38699 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5134)
      %38700 = stablehlo.reduce(%38698 init: %38699) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5134)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5134)
      } loc(#loc5134)
      %38701 = stablehlo.reshape %38700 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5135)
      %38702 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5136)
      %38703 = stablehlo.reduce(%38701 init: %38702) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5136)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5136)
      } loc(#loc5136)
      %38704 = stablehlo.broadcast_in_dim %38703, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5137)
      %38705 = func.call @_where_1117(%30501#1, %38704) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38706 = arith.select %ptrue, %38698, %38698 : tensor<1x1x2048x2048xbf16>
      %38707:2 = func.call @_einsum_1118(%30502#0, %30933, %38706) : (tensor<1x8192x2048xbf16>, tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) loc(#loc5138)
      %38708 = func.call @_where_1119(%30502#1, %38707#1) : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc4056)
      %38709 = stablehlo.broadcast_in_dim %30931, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2137)
      %38710 = stablehlo.multiply %38707#0, %38709 : tensor<1x1x2048x8192xbf16> loc(#loc2137)
      %38711 = stablehlo.multiply %30902, %38710 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38712 = stablehlo.constant dense<5.000000e-01> : tensor<bf16> loc(#loc1984)
      %38713 = stablehlo.broadcast_in_dim %38712, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38714 = stablehlo.multiply %38713, %38711 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38715 = stablehlo.multiply %38714, %30919 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38716 = stablehlo.multiply %38715, %30916 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38717 = stablehlo.add %38715, %38716 : tensor<1x1x2048x8192xbf16> loc(#loc5139)
      %38718 = stablehlo.constant dense<7.968750e-01> : tensor<bf16> loc(#loc1984)
      %38719 = stablehlo.broadcast_in_dim %38718, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38720 = stablehlo.multiply %38719, %38717 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38721 = stablehlo.constant dense<4.467770e-02> : tensor<bf16> loc(#loc1984)
      %38722 = stablehlo.broadcast_in_dim %38721, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38723 = stablehlo.multiply %38722, %38720 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38724 = stablehlo.multiply %38723, %30908 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38725 = stablehlo.add %38720, %38724 : tensor<1x1x2048x8192xbf16> loc(#loc5139)
      %38726 = stablehlo.multiply %38710, %30925 : tensor<1x1x2048x8192xbf16> loc(#loc2130)
      %38727 = stablehlo.add %38725, %38726 : tensor<1x1x2048x8192xbf16> loc(#loc5139)
      %38728 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5140)
      %38729 = stablehlo.reduce(%38727 init: %38728) across dimensions = [2] : (tensor<1x1x2048x8192xbf16>, tensor<bf16>) -> tensor<1x1x8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5140)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5140)
      } loc(#loc5140)
      %38730 = stablehlo.reshape %38729 : (tensor<1x1x8192xbf16>) -> tensor<1x1x1x8192xbf16> loc(#loc5141)
      %38731 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5142)
      %38732 = stablehlo.reduce(%38730 init: %38731) across dimensions = [0, 1, 2] : (tensor<1x1x1x8192xbf16>, tensor<bf16>) -> tensor<8192xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5142)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5142)
      } loc(#loc5142)
      %38733 = stablehlo.broadcast_in_dim %38732, dims = [1] : (tensor<8192xbf16>) -> tensor<1x8192xbf16> loc(#loc5143)
      %38734 = func.call @_where_1120(%30499#1, %38733) : (tensor<1x8192xi1>, tensor<1x8192xbf16>) -> tensor<1x8192xbf16> loc(#loc4056)
      %38735 = arith.select %ptrue, %38727, %38727 : tensor<1x1x2048x8192xbf16>
      %38736:2 = func.call @_einsum_1121(%30500#0, %30897, %38735) : (tensor<1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) loc(#loc5144)
      %38737 = func.call @_where_1122(%30500#1, %38736#1) : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc4056)
      %38738 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5145)
      %38739 = stablehlo.reduce(%38736#0 init: %38738) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5145)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5145)
      } loc(#loc5145)
      %38740 = stablehlo.reshape %38739 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5146)
      %38741 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5147)
      %38742 = stablehlo.reduce(%38740 init: %38741) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5147)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5147)
      } loc(#loc5147)
      %38743 = stablehlo.broadcast_in_dim %38742, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5148)
      %38744 = func.call @_where_1123(%30503#1, %38743) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38745 = stablehlo.multiply %30888, %38736#0 : tensor<1x1x2048x2048xbf16> loc(#loc2123)
      %38746 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5149)
      %38747 = stablehlo.reduce(%38745 init: %38746) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5149)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5149)
      } loc(#loc5149)
      %38748 = stablehlo.reshape %38747 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5150)
      %38749 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5151)
      %38750 = stablehlo.reduce(%38748 init: %38749) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5151)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5151)
      } loc(#loc5151)
      %38751 = stablehlo.broadcast_in_dim %38750, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5152)
      %38752 = func.call @_where_1124(%30504#1, %38751) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38753 = stablehlo.broadcast_in_dim %30892, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2123)
      %38754 = stablehlo.multiply %38736#0, %38753 : tensor<1x1x2048x2048xbf16> loc(#loc2123)
      %38755 = stablehlo.multiply %30878, %38754 : tensor<1x1x2048x2048xbf16> loc(#loc2120)
      %38756 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5153)
      %38757 = stablehlo.reduce(%38755 init: %38756) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5153)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5153)
      } loc(#loc5153)
      %38758 = stablehlo.reshape %38757 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5154)
      %38759 = stablehlo.multiply %38758, %30886 : tensor<1x1x2048x1xbf16> loc(#loc2120)
      %38760 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38761 = stablehlo.broadcast_in_dim %38760, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2114)
      %38762 = stablehlo.divide %38759, %38761 : tensor<1x1x2048x1xbf16> loc(#loc2114)
      %38763 = stablehlo.convert %38762 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2110)
      %38764 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5155)
      %38765 = stablehlo.reduce(%38763 init: %38764) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5155)
        stablehlo.return %39096 : tensor<f32> loc(#loc5155)
      } loc(#loc5155)
      %38766 = stablehlo.broadcast_in_dim %38765, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5156)
      %38767 = stablehlo.broadcast_in_dim %38766, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5157)
      %38768 = stablehlo.convert %38767 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2113)
      %38769 = stablehlo.multiply %38768, %30868 : tensor<1x1x2048x2048xbf16> loc(#loc2116)
      %38770 = stablehlo.negate %38769 : tensor<1x1x2048x2048xbf16> loc(#loc5158)
      %38771 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5159)
      %38772 = stablehlo.reduce(%38770 init: %38771) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5159)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5159)
      } loc(#loc5159)
      %38773 = stablehlo.reshape %38772 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5160)
      %38774 = stablehlo.add %38693, %38769 : tensor<1x1x2048x2048xbf16> loc(#loc5161)
      %38775 = stablehlo.broadcast_in_dim %30882, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2120)
      %38776 = stablehlo.multiply %38754, %38775 : tensor<1x1x2048x2048xbf16> loc(#loc2120)
      %38777 = stablehlo.negate %38776 : tensor<1x1x2048x2048xbf16> loc(#loc5158)
      %38778 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5159)
      %38779 = stablehlo.reduce(%38777 init: %38778) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5159)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5159)
      } loc(#loc5159)
      %38780 = stablehlo.reshape %38779 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5160)
      %38781 = stablehlo.add %38774, %38776 : tensor<1x1x2048x2048xbf16> loc(#loc5161)
      %38782 = stablehlo.add %38773, %38780 : tensor<1x1x2048x1xbf16> loc(#loc5161)
      %38783 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38784 = stablehlo.broadcast_in_dim %38783, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2114)
      %38785 = stablehlo.divide %38782, %38784 : tensor<1x1x2048x1xbf16> loc(#loc2114)
      %38786 = stablehlo.convert %38785 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2110)
      %38787 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5155)
      %38788 = stablehlo.reduce(%38786 init: %38787) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5155)
        stablehlo.return %39096 : tensor<f32> loc(#loc5155)
      } loc(#loc5155)
      %38789 = stablehlo.broadcast_in_dim %38788, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5156)
      %38790 = stablehlo.broadcast_in_dim %38789, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5157)
      %38791 = stablehlo.convert %38790 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2113)
      %38792 = stablehlo.add %38781, %38791 : tensor<1x1x2048x2048xbf16> loc(#loc5162)
      %38793 = stablehlo.add %35419, %38792 : tensor<1x1x2048x2048xbf16> loc(#loc5163)
      %38794 = arith.select %ptrue, %38792, %38792 : tensor<1x1x2048x2048xbf16>
      %38795:2 = func.call @_einsum_1125(%30509#0, %30850, %38794) : (tensor<1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) loc(#loc5164)
      %38796 = func.call @_where_1126(%30509#1, %38795#1) : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc4056)
      %38797 = arith.select %ptrue, %38795#0, %38795#0 : tensor<1x1x2048x32x64xbf16>
      %38798:2 = func.call @_einsum_1127(%30848, %30805, %38797) : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc5165)
      %38799 = arith.select %ptrue, %38798#1, %38798#1 : tensor<1x1x2048x32x64xbf16>
      %38800 = stablehlo.broadcast_in_dim %38799, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5166)
      %38801 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38802 = stablehlo.pad %38800, %38801, low = [0, 2, 0, 0, 0, 0], high = [0, 0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5167)
      %38803 = stablehlo.convert %38798#0 : (tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5168)
      %38804 = stablehlo.broadcast_in_dim %30847, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2101)
      %38805 = stablehlo.multiply %38803, %38804 : tensor<1x1x32x2048x2048xf32> loc(#loc2101)
      %38806 = stablehlo.multiply %38805, %30838 : tensor<1x1x32x2048x2048xf32> loc(#loc2101)
      %38807 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc2099)
      %38808 = stablehlo.reduce(%38806 init: %38807) across dimensions = [4] : (tensor<1x1x32x2048x2048xf32>, tensor<f32>) -> tensor<1x1x32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc2099)
        stablehlo.return %39096 : tensor<f32> loc(#loc2099)
      } loc(#loc2099)
      %38809 = stablehlo.reshape %38808 : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x1xf32> loc(#loc5169)
      %38810 = stablehlo.negate %38809 : tensor<1x1x32x2048x1xf32> loc(#loc5170)
      %38811 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5171)
      %38812 = stablehlo.reduce(%38810 init: %38811) across dimensions = [0, 1, 4] : (tensor<1x1x32x2048x1xf32>, tensor<f32>) -> tensor<32x2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5171)
        stablehlo.return %39096 : tensor<f32> loc(#loc5171)
      } loc(#loc5171)
      %38813 = stablehlo.broadcast_in_dim %38812, dims = [2, 3] : (tensor<32x2048xf32>) -> tensor<1x1x32x2048xf32> loc(#loc5172)
      %38814 = stablehlo.broadcast_in_dim %38813, dims = [0, 1, 2, 3] : (tensor<1x1x32x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5173)
      %38815 = stablehlo.broadcast_in_dim %30841, dims = [0, 1, 2, 3, 4] : (tensor<1x1x32x2048x1xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc2100)
      %38816 = stablehlo.divide %38803, %38815 : tensor<1x1x32x2048x2048xf32> loc(#loc2100)
      %38817 = stablehlo.add %38814, %38816 : tensor<1x1x32x2048x2048xf32> loc(#loc5174)
      %38818 = stablehlo.multiply %38817, %30838 : tensor<1x1x32x2048x2048xf32> loc(#loc2101)
      %38819 = func.call @_where_1128(%30832#1, %38818) : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5175)
      %38820 = stablehlo.convert %38819 : (tensor<1x1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5176)
      %38821 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38822 = stablehlo.broadcast_in_dim %38821, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2089)
      %38823 = stablehlo.multiply %38822, %38820 : tensor<1x1x32x2048x2048xbf16> loc(#loc2089)
      %38824 = stablehlo.multiply %38823, %30820 : tensor<1x1x32x2048x2048xbf16> loc(#loc2089)
      %38825 = stablehlo.multiply %38824, %30817 : tensor<1x1x32x2048x2048xbf16> loc(#loc2089)
      %38826 = stablehlo.add %38824, %38825 : tensor<1x1x32x2048x2048xbf16> loc(#loc5177)
      %38827 = stablehlo.constant dense<5.000000e+01> : tensor<bf16> loc(#loc1984)
      %38828 = stablehlo.broadcast_in_dim %38827, dims = [] : (tensor<bf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc2086)
      %38829 = stablehlo.divide %38826, %38828 : tensor<1x1x32x2048x2048xbf16> loc(#loc2086)
      %38830:2 = func.call @_einsum_1129(%30804, %30812, %38829) : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) loc(#loc5178)
      %38831 = arith.select %ptrue, %38830#1, %38830#1 : tensor<1x1x2048x32x64xbf16>
      %38832 = stablehlo.broadcast_in_dim %38831, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5166)
      %38833 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38834 = stablehlo.pad %38832, %38833, low = [0, 1, 0, 0, 0, 0], high = [0, 1, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5179)
      %38835 = stablehlo.add %38802, %38834 : tensor<1x3x1x2048x32x64xbf16> loc(#loc5180)
      %38836 = stablehlo.multiply %30803, %38830#0 : tensor<1x1x2048x32x64xbf16> loc(#loc2084)
      %38837 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5181)
      %38838 = stablehlo.reduce(%38836 init: %38837) across dimensions = [2, 3] : (tensor<1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x1x64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5181)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5181)
      } loc(#loc5181)
      %38839 = stablehlo.reshape %38838 : (tensor<1x1x64xbf16>) -> tensor<1x1x1x1x64xbf16> loc(#loc5182)
      %38840 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5183)
      %38841 = stablehlo.reduce(%38839 init: %38840) across dimensions = [0, 1, 2, 3] : (tensor<1x1x1x1x64xbf16>, tensor<bf16>) -> tensor<64xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5183)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5183)
      } loc(#loc5183)
      %38842 = stablehlo.broadcast_in_dim %38841, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5184)
      %38843 = stablehlo.constant dense<1.806640e-01> : tensor<bf16> loc(#loc1984)
      %38844 = stablehlo.broadcast_in_dim %38843, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc2082)
      %38845 = stablehlo.multiply %38844, %38842 : tensor<1x64xbf16> loc(#loc2082)
      %38846 = func.call @softplus_1130(%30806#1, %38845) : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc5185)
      %38847 = func.call @_where_1131(%30508#1, %38846) : (tensor<1x64xi1>, tensor<1x64xbf16>) -> tensor<1x64xbf16> loc(#loc4056)
      %38848 = stablehlo.broadcast_in_dim %30810, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc2084)
      %38849 = stablehlo.multiply %38830#0, %38848 : tensor<1x1x2048x32x64xbf16> loc(#loc2084)
      %38850 = arith.select %ptrue, %38849, %38849 : tensor<1x1x2048x32x64xbf16>
      %38851 = stablehlo.broadcast_in_dim %38850, dims = [0, 2, 3, 4, 5] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x1x2048x32x64xbf16> loc(#loc5166)
      %38852 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38853 = stablehlo.pad %38851, %38852, low = [0, 0, 0, 0, 0, 0], high = [0, 2, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0, 0] : (tensor<1x1x1x2048x32x64xbf16>, tensor<bf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5186)
      %38854 = stablehlo.add %38835, %38853 : tensor<1x3x1x2048x32x64xbf16> loc(#loc5180)
      %38855:2 = func.call @_einsum_1132(%30795, %30507#0, %38854) : (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x3x1x2048x32x64xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) loc(#loc5187)
      %38856 = func.call @_where_1133(%30507#1, %38855#1) : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc4056)
      %38857 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5188)
      %38858 = stablehlo.reduce(%38855#0 init: %38857) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5188)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5188)
      } loc(#loc5188)
      %38859 = stablehlo.reshape %38858 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5189)
      %38860 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5190)
      %38861 = stablehlo.reduce(%38859 init: %38860) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5190)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5190)
      } loc(#loc5190)
      %38862 = stablehlo.broadcast_in_dim %38861, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5191)
      %38863 = func.call @_where_1134(%30505#1, %38862) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38864 = stablehlo.multiply %30786, %38855#0 : tensor<1x1x2048x2048xbf16> loc(#loc2072)
      %38865 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5192)
      %38866 = stablehlo.reduce(%38864 init: %38865) across dimensions = [2] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5192)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5192)
      } loc(#loc5192)
      %38867 = stablehlo.reshape %38866 : (tensor<1x1x2048xbf16>) -> tensor<1x1x1x2048xbf16> loc(#loc5193)
      %38868 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5194)
      %38869 = stablehlo.reduce(%38867 init: %38868) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048xbf16>, tensor<bf16>) -> tensor<2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5194)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5194)
      } loc(#loc5194)
      %38870 = stablehlo.broadcast_in_dim %38869, dims = [1] : (tensor<2048xbf16>) -> tensor<1x2048xbf16> loc(#loc5195)
      %38871 = func.call @_where_1135(%30506#1, %38870) : (tensor<1x2048xi1>, tensor<1x2048xbf16>) -> tensor<1x2048xbf16> loc(#loc4056)
      %38872 = stablehlo.broadcast_in_dim %30790, dims = [0, 1, 2, 3] : (tensor<1x1x1x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2072)
      %38873 = stablehlo.multiply %38855#0, %38872 : tensor<1x1x2048x2048xbf16> loc(#loc2072)
      %38874 = stablehlo.multiply %30776, %38873 : tensor<1x1x2048x2048xbf16> loc(#loc2069)
      %38875 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5196)
      %38876 = stablehlo.reduce(%38874 init: %38875) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5196)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5196)
      } loc(#loc5196)
      %38877 = stablehlo.reshape %38876 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5197)
      %38878 = stablehlo.multiply %38877, %30784 : tensor<1x1x2048x1xbf16> loc(#loc2069)
      %38879 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38880 = stablehlo.broadcast_in_dim %38879, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2063)
      %38881 = stablehlo.divide %38878, %38880 : tensor<1x1x2048x1xbf16> loc(#loc2063)
      %38882 = stablehlo.convert %38881 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2059)
      %38883 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5198)
      %38884 = stablehlo.reduce(%38882 init: %38883) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5198)
        stablehlo.return %39096 : tensor<f32> loc(#loc5198)
      } loc(#loc5198)
      %38885 = stablehlo.broadcast_in_dim %38884, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5199)
      %38886 = stablehlo.broadcast_in_dim %38885, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5200)
      %38887 = stablehlo.convert %38886 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2062)
      %38888 = stablehlo.multiply %38887, %30766 : tensor<1x1x2048x2048xbf16> loc(#loc2065)
      %38889 = stablehlo.negate %38888 : tensor<1x1x2048x2048xbf16> loc(#loc5201)
      %38890 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5202)
      %38891 = stablehlo.reduce(%38889 init: %38890) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5202)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5202)
      } loc(#loc5202)
      %38892 = stablehlo.reshape %38891 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5203)
      %38893 = stablehlo.add %38793, %38888 : tensor<1x1x2048x2048xbf16> loc(#loc5204)
      %38894 = stablehlo.broadcast_in_dim %30780, dims = [0, 1, 2, 3] : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc2069)
      %38895 = stablehlo.multiply %38873, %38894 : tensor<1x1x2048x2048xbf16> loc(#loc2069)
      %38896 = stablehlo.negate %38895 : tensor<1x1x2048x2048xbf16> loc(#loc5201)
      %38897 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5202)
      %38898 = stablehlo.reduce(%38896 init: %38897) across dimensions = [3] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048xbf16>
       reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5202)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5202)
      } loc(#loc5202)
      %38899 = stablehlo.reshape %38898 : (tensor<1x1x2048xbf16>) -> tensor<1x1x2048x1xbf16> loc(#loc5203)
      %38900 = stablehlo.add %38893, %38895 : tensor<1x1x2048x2048xbf16> loc(#loc5204)
      %38901 = stablehlo.add %38892, %38899 : tensor<1x1x2048x1xbf16> loc(#loc5204)
      %38902 = stablehlo.constant dense<2.048000e+03> : tensor<bf16> loc(#loc1984)
      %38903 = stablehlo.broadcast_in_dim %38902, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x1xbf16> loc(#loc2063)
      %38904 = stablehlo.divide %38901, %38903 : tensor<1x1x2048x1xbf16> loc(#loc2063)
      %38905 = stablehlo.convert %38904 : (tensor<1x1x2048x1xbf16>) -> tensor<1x1x2048x1xf32> loc(#loc2059)
      %38906 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5198)
      %38907 = stablehlo.reduce(%38905 init: %38906) across dimensions = [0, 1, 3] : (tensor<1x1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
       reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
        %39096 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5198)
        stablehlo.return %39096 : tensor<f32> loc(#loc5198)
      } loc(#loc5198)
      %38908 = stablehlo.broadcast_in_dim %38907, dims = [2] : (tensor<2048xf32>) -> tensor<1x1x2048xf32> loc(#loc5199)
      %38909 = stablehlo.broadcast_in_dim %38908, dims = [0, 1, 2] : (tensor<1x1x2048xf32>) -> tensor<1x1x2048x2048xf32> loc(#loc5200)
      %38910 = stablehlo.convert %38909 : (tensor<1x1x2048x2048xf32>) -> tensor<1x1x2048x2048xbf16> loc(#loc2062)
      %38911 = stablehlo.add %38900, %38910 : tensor<1x1x2048x2048xbf16> loc(#loc5205)
      %38912 = arith.select %ptrue, %38911, %38911 : tensor<1x1x2048x2048xbf16>
      %38913:2 = func.call @_where_1136(%30417, %38912) : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xbf16>) -> (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x2048xbf16>) loc(#loc5206)
      %38914 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc1984)
      %38915 = stablehlo.broadcast_in_dim %38914, dims = [] : (tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5207)
      %38916 = "stablehlo.scatter"(%38915, %30412, %38913#0) ({
      ^bb0(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown)):
        %39096 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5209)
        stablehlo.return %39096 : tensor<bf16> loc(#loc5208)
      }) {indices_are_sorted = true, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0, 1, 2, 3], inserted_window_dims = [0], scatter_dims_to_operand_dims = [0]>, unique_indices = true} : (tensor<1x1x1x2048x2048xbf16>, tensor<1xi32>, tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5208)
      %38917 = stablehlo.add %iterArg_229, %38916 : tensor<1x1x1x2048x2048xbf16> loc(#loc5210)
      %38918 = stablehlo.add %iterArg_230, %38734 : tensor<1x8192xbf16> loc(#loc5210)
      %38919 = stablehlo.add %iterArg_231, %38737 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38920 = stablehlo.add %iterArg_232, %38705 : tensor<1x2048xbf16> loc(#loc5210)
      %38921 = stablehlo.add %iterArg_233, %38708 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38922 = stablehlo.add %iterArg_234, %38744 : tensor<1x2048xbf16> loc(#loc5210)
      %38923 = stablehlo.add %iterArg_235, %38752 : tensor<1x2048xbf16> loc(#loc5210)
      %38924 = stablehlo.add %iterArg_236, %38863 : tensor<1x2048xbf16> loc(#loc5210)
      %38925 = stablehlo.add %iterArg_237, %38871 : tensor<1x2048xbf16> loc(#loc5210)
      %38926 = stablehlo.add %iterArg_238, %38856 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38927 = stablehlo.add %iterArg_239, %38847 : tensor<1x64xbf16> loc(#loc5210)
      %38928 = stablehlo.add %iterArg_240, %38796 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38929 = stablehlo.add %iterArg_241, %38516 : tensor<1x8192xbf16> loc(#loc5210)
      %38930 = stablehlo.add %iterArg_242, %38519 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38931 = stablehlo.add %iterArg_243, %38487 : tensor<1x2048xbf16> loc(#loc5210)
      %38932 = stablehlo.add %iterArg_244, %38490 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38933 = stablehlo.add %iterArg_245, %38526 : tensor<1x2048xbf16> loc(#loc5210)
      %38934 = stablehlo.add %iterArg_246, %38534 : tensor<1x2048xbf16> loc(#loc5210)
      %38935 = stablehlo.add %iterArg_247, %38645 : tensor<1x2048xbf16> loc(#loc5210)
      %38936 = stablehlo.add %iterArg_248, %38653 : tensor<1x2048xbf16> loc(#loc5210)
      %38937 = stablehlo.add %iterArg_249, %38638 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38938 = stablehlo.add %iterArg_250, %38629 : tensor<1x64xbf16> loc(#loc5210)
      %38939 = stablehlo.add %iterArg_251, %38578 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38940 = stablehlo.add %iterArg_252, %36554 : tensor<1x8192xbf16> loc(#loc5210)
      %38941 = stablehlo.add %iterArg_253, %36557 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38942 = stablehlo.add %iterArg_254, %36525 : tensor<1x2048xbf16> loc(#loc5210)
      %38943 = stablehlo.add %iterArg_255, %36528 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38944 = stablehlo.add %iterArg_256, %36564 : tensor<1x2048xbf16> loc(#loc5210)
      %38945 = stablehlo.add %iterArg_257, %36572 : tensor<1x2048xbf16> loc(#loc5210)
      %38946 = stablehlo.add %iterArg_258, %36683 : tensor<1x2048xbf16> loc(#loc5210)
      %38947 = stablehlo.add %iterArg_259, %36691 : tensor<1x2048xbf16> loc(#loc5210)
      %38948 = stablehlo.add %iterArg_260, %36676 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38949 = stablehlo.add %iterArg_261, %36667 : tensor<1x64xbf16> loc(#loc5210)
      %38950 = stablehlo.add %iterArg_262, %36616 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38951 = stablehlo.add %iterArg_263, %36336 : tensor<1x8192xbf16> loc(#loc5210)
      %38952 = stablehlo.add %iterArg_264, %36339 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38953 = stablehlo.add %iterArg_265, %36307 : tensor<1x2048xbf16> loc(#loc5210)
      %38954 = stablehlo.add %iterArg_266, %36310 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38955 = stablehlo.add %iterArg_267, %36346 : tensor<1x2048xbf16> loc(#loc5210)
      %38956 = stablehlo.add %iterArg_268, %36354 : tensor<1x2048xbf16> loc(#loc5210)
      %38957 = stablehlo.add %iterArg_269, %36465 : tensor<1x2048xbf16> loc(#loc5210)
      %38958 = stablehlo.add %iterArg_270, %36473 : tensor<1x2048xbf16> loc(#loc5210)
      %38959 = stablehlo.add %iterArg_271, %36458 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38960 = stablehlo.add %iterArg_272, %36449 : tensor<1x64xbf16> loc(#loc5210)
      %38961 = stablehlo.add %iterArg_273, %36398 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38962 = stablehlo.add %iterArg_274, %36118 : tensor<1x8192xbf16> loc(#loc5210)
      %38963 = stablehlo.add %iterArg_275, %36121 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38964 = stablehlo.add %iterArg_276, %36089 : tensor<1x2048xbf16> loc(#loc5210)
      %38965 = stablehlo.add %iterArg_277, %36092 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38966 = stablehlo.add %iterArg_278, %36128 : tensor<1x2048xbf16> loc(#loc5210)
      %38967 = stablehlo.add %iterArg_279, %36136 : tensor<1x2048xbf16> loc(#loc5210)
      %38968 = stablehlo.add %iterArg_280, %36247 : tensor<1x2048xbf16> loc(#loc5210)
      %38969 = stablehlo.add %iterArg_281, %36255 : tensor<1x2048xbf16> loc(#loc5210)
      %38970 = stablehlo.add %iterArg_282, %36240 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38971 = stablehlo.add %iterArg_283, %36231 : tensor<1x64xbf16> loc(#loc5210)
      %38972 = stablehlo.add %iterArg_284, %36180 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38973 = stablehlo.add %iterArg_285, %35900 : tensor<1x8192xbf16> loc(#loc5210)
      %38974 = stablehlo.add %iterArg_286, %35903 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38975 = stablehlo.add %iterArg_287, %35871 : tensor<1x2048xbf16> loc(#loc5210)
      %38976 = stablehlo.add %iterArg_288, %35874 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38977 = stablehlo.add %iterArg_289, %35910 : tensor<1x2048xbf16> loc(#loc5210)
      %38978 = stablehlo.add %iterArg_290, %35918 : tensor<1x2048xbf16> loc(#loc5210)
      %38979 = stablehlo.add %iterArg_291, %36029 : tensor<1x2048xbf16> loc(#loc5210)
      %38980 = stablehlo.add %iterArg_292, %36037 : tensor<1x2048xbf16> loc(#loc5210)
      %38981 = stablehlo.add %iterArg_293, %36022 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38982 = stablehlo.add %iterArg_294, %36013 : tensor<1x64xbf16> loc(#loc5210)
      %38983 = stablehlo.add %iterArg_295, %35962 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38984 = stablehlo.add %iterArg_296, %35682 : tensor<1x8192xbf16> loc(#loc5210)
      %38985 = stablehlo.add %iterArg_297, %35685 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38986 = stablehlo.add %iterArg_298, %35653 : tensor<1x2048xbf16> loc(#loc5210)
      %38987 = stablehlo.add %iterArg_299, %35656 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38988 = stablehlo.add %iterArg_300, %35692 : tensor<1x2048xbf16> loc(#loc5210)
      %38989 = stablehlo.add %iterArg_301, %35700 : tensor<1x2048xbf16> loc(#loc5210)
      %38990 = stablehlo.add %iterArg_302, %35811 : tensor<1x2048xbf16> loc(#loc5210)
      %38991 = stablehlo.add %iterArg_303, %35819 : tensor<1x2048xbf16> loc(#loc5210)
      %38992 = stablehlo.add %iterArg_304, %35804 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %38993 = stablehlo.add %iterArg_305, %35795 : tensor<1x64xbf16> loc(#loc5210)
      %38994 = stablehlo.add %iterArg_306, %35744 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %38995 = stablehlo.add %iterArg_307, %35464 : tensor<1x8192xbf16> loc(#loc5210)
      %38996 = stablehlo.add %iterArg_308, %35467 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %38997 = stablehlo.add %iterArg_309, %35435 : tensor<1x2048xbf16> loc(#loc5210)
      %38998 = stablehlo.add %iterArg_310, %35438 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %38999 = stablehlo.add %iterArg_311, %35474 : tensor<1x2048xbf16> loc(#loc5210)
      %39000 = stablehlo.add %iterArg_312, %35482 : tensor<1x2048xbf16> loc(#loc5210)
      %39001 = stablehlo.add %iterArg_313, %35593 : tensor<1x2048xbf16> loc(#loc5210)
      %39002 = stablehlo.add %iterArg_314, %35601 : tensor<1x2048xbf16> loc(#loc5210)
      %39003 = stablehlo.add %iterArg_315, %35586 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39004 = stablehlo.add %iterArg_316, %35577 : tensor<1x64xbf16> loc(#loc5210)
      %39005 = stablehlo.add %iterArg_317, %35526 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39006 = stablehlo.add %iterArg_318, %38298 : tensor<1x8192xbf16> loc(#loc5210)
      %39007 = stablehlo.add %iterArg_319, %38301 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39008 = stablehlo.add %iterArg_320, %38269 : tensor<1x2048xbf16> loc(#loc5210)
      %39009 = stablehlo.add %iterArg_321, %38272 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39010 = stablehlo.add %iterArg_322, %38308 : tensor<1x2048xbf16> loc(#loc5210)
      %39011 = stablehlo.add %iterArg_323, %38316 : tensor<1x2048xbf16> loc(#loc5210)
      %39012 = stablehlo.add %iterArg_324, %38427 : tensor<1x2048xbf16> loc(#loc5210)
      %39013 = stablehlo.add %iterArg_325, %38435 : tensor<1x2048xbf16> loc(#loc5210)
      %39014 = stablehlo.add %iterArg_326, %38420 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39015 = stablehlo.add %iterArg_327, %38411 : tensor<1x64xbf16> loc(#loc5210)
      %39016 = stablehlo.add %iterArg_328, %38360 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39017 = stablehlo.add %iterArg_329, %38080 : tensor<1x8192xbf16> loc(#loc5210)
      %39018 = stablehlo.add %iterArg_330, %38083 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39019 = stablehlo.add %iterArg_331, %38051 : tensor<1x2048xbf16> loc(#loc5210)
      %39020 = stablehlo.add %iterArg_332, %38054 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39021 = stablehlo.add %iterArg_333, %38090 : tensor<1x2048xbf16> loc(#loc5210)
      %39022 = stablehlo.add %iterArg_334, %38098 : tensor<1x2048xbf16> loc(#loc5210)
      %39023 = stablehlo.add %iterArg_335, %38209 : tensor<1x2048xbf16> loc(#loc5210)
      %39024 = stablehlo.add %iterArg_336, %38217 : tensor<1x2048xbf16> loc(#loc5210)
      %39025 = stablehlo.add %iterArg_337, %38202 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39026 = stablehlo.add %iterArg_338, %38193 : tensor<1x64xbf16> loc(#loc5210)
      %39027 = stablehlo.add %iterArg_339, %38142 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39028 = stablehlo.add %iterArg_340, %37862 : tensor<1x8192xbf16> loc(#loc5210)
      %39029 = stablehlo.add %iterArg_341, %37865 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39030 = stablehlo.add %iterArg_342, %37833 : tensor<1x2048xbf16> loc(#loc5210)
      %39031 = stablehlo.add %iterArg_343, %37836 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39032 = stablehlo.add %iterArg_344, %37872 : tensor<1x2048xbf16> loc(#loc5210)
      %39033 = stablehlo.add %iterArg_345, %37880 : tensor<1x2048xbf16> loc(#loc5210)
      %39034 = stablehlo.add %iterArg_346, %37991 : tensor<1x2048xbf16> loc(#loc5210)
      %39035 = stablehlo.add %iterArg_347, %37999 : tensor<1x2048xbf16> loc(#loc5210)
      %39036 = stablehlo.add %iterArg_348, %37984 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39037 = stablehlo.add %iterArg_349, %37975 : tensor<1x64xbf16> loc(#loc5210)
      %39038 = stablehlo.add %iterArg_350, %37924 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39039 = stablehlo.add %iterArg_351, %37644 : tensor<1x8192xbf16> loc(#loc5210)
      %39040 = stablehlo.add %iterArg_352, %37647 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39041 = stablehlo.add %iterArg_353, %37615 : tensor<1x2048xbf16> loc(#loc5210)
      %39042 = stablehlo.add %iterArg_354, %37618 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39043 = stablehlo.add %iterArg_355, %37654 : tensor<1x2048xbf16> loc(#loc5210)
      %39044 = stablehlo.add %iterArg_356, %37662 : tensor<1x2048xbf16> loc(#loc5210)
      %39045 = stablehlo.add %iterArg_357, %37773 : tensor<1x2048xbf16> loc(#loc5210)
      %39046 = stablehlo.add %iterArg_358, %37781 : tensor<1x2048xbf16> loc(#loc5210)
      %39047 = stablehlo.add %iterArg_359, %37766 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39048 = stablehlo.add %iterArg_360, %37757 : tensor<1x64xbf16> loc(#loc5210)
      %39049 = stablehlo.add %iterArg_361, %37706 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39050 = stablehlo.add %iterArg_362, %37426 : tensor<1x8192xbf16> loc(#loc5210)
      %39051 = stablehlo.add %iterArg_363, %37429 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39052 = stablehlo.add %iterArg_364, %37397 : tensor<1x2048xbf16> loc(#loc5210)
      %39053 = stablehlo.add %iterArg_365, %37400 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39054 = stablehlo.add %iterArg_366, %37436 : tensor<1x2048xbf16> loc(#loc5210)
      %39055 = stablehlo.add %iterArg_367, %37444 : tensor<1x2048xbf16> loc(#loc5210)
      %39056 = stablehlo.add %iterArg_368, %37555 : tensor<1x2048xbf16> loc(#loc5210)
      %39057 = stablehlo.add %iterArg_369, %37563 : tensor<1x2048xbf16> loc(#loc5210)
      %39058 = stablehlo.add %iterArg_370, %37548 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39059 = stablehlo.add %iterArg_371, %37539 : tensor<1x64xbf16> loc(#loc5210)
      %39060 = stablehlo.add %iterArg_372, %37488 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39061 = stablehlo.add %iterArg_373, %37208 : tensor<1x8192xbf16> loc(#loc5210)
      %39062 = stablehlo.add %iterArg_374, %37211 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39063 = stablehlo.add %iterArg_375, %37179 : tensor<1x2048xbf16> loc(#loc5210)
      %39064 = stablehlo.add %iterArg_376, %37182 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39065 = stablehlo.add %iterArg_377, %37218 : tensor<1x2048xbf16> loc(#loc5210)
      %39066 = stablehlo.add %iterArg_378, %37226 : tensor<1x2048xbf16> loc(#loc5210)
      %39067 = stablehlo.add %iterArg_379, %37337 : tensor<1x2048xbf16> loc(#loc5210)
      %39068 = stablehlo.add %iterArg_380, %37345 : tensor<1x2048xbf16> loc(#loc5210)
      %39069 = stablehlo.add %iterArg_381, %37330 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39070 = stablehlo.add %iterArg_382, %37321 : tensor<1x64xbf16> loc(#loc5210)
      %39071 = stablehlo.add %iterArg_383, %37270 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39072 = stablehlo.add %iterArg_384, %36990 : tensor<1x8192xbf16> loc(#loc5210)
      %39073 = stablehlo.add %iterArg_385, %36993 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39074 = stablehlo.add %iterArg_386, %36961 : tensor<1x2048xbf16> loc(#loc5210)
      %39075 = stablehlo.add %iterArg_387, %36964 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39076 = stablehlo.add %iterArg_388, %37000 : tensor<1x2048xbf16> loc(#loc5210)
      %39077 = stablehlo.add %iterArg_389, %37008 : tensor<1x2048xbf16> loc(#loc5210)
      %39078 = stablehlo.add %iterArg_390, %37119 : tensor<1x2048xbf16> loc(#loc5210)
      %39079 = stablehlo.add %iterArg_391, %37127 : tensor<1x2048xbf16> loc(#loc5210)
      %39080 = stablehlo.add %iterArg_392, %37112 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39081 = stablehlo.add %iterArg_393, %37103 : tensor<1x64xbf16> loc(#loc5210)
      %39082 = stablehlo.add %iterArg_394, %37052 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39083 = stablehlo.add %iterArg_395, %36772 : tensor<1x8192xbf16> loc(#loc5210)
      %39084 = stablehlo.add %iterArg_396, %36775 : tensor<1x2048x8192xbf16> loc(#loc5210)
      %39085 = stablehlo.add %iterArg_397, %36743 : tensor<1x2048xbf16> loc(#loc5210)
      %39086 = stablehlo.add %iterArg_398, %36746 : tensor<1x8192x2048xbf16> loc(#loc5210)
      %39087 = stablehlo.add %iterArg_399, %36782 : tensor<1x2048xbf16> loc(#loc5210)
      %39088 = stablehlo.add %iterArg_400, %36790 : tensor<1x2048xbf16> loc(#loc5210)
      %39089 = stablehlo.add %iterArg_401, %36901 : tensor<1x2048xbf16> loc(#loc5210)
      %39090 = stablehlo.add %iterArg_402, %36909 : tensor<1x2048xbf16> loc(#loc5210)
      %39091 = stablehlo.add %iterArg_403, %36894 : tensor<1x3x2048x32x64xbf16> loc(#loc5210)
      %39092 = stablehlo.add %iterArg_404, %36885 : tensor<1x64xbf16> loc(#loc5210)
      %39093 = stablehlo.add %iterArg_405, %36834 : tensor<1x2048x32x64xbf16> loc(#loc5210)
      %39094 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1971)
      %39095 = stablehlo.add %iterArg_228, %39094 : tensor<i32> loc(#loc1976)
      stablehlo.return %iterArg, %iterArg_0, %iterArg_1, %iterArg_2, %iterArg_3, %iterArg_4, %iterArg_5, %iterArg_6, %iterArg_7, %iterArg_8, %iterArg_9, %iterArg_10, %iterArg_11, %iterArg_12, %iterArg_13, %iterArg_14, %iterArg_15, %iterArg_16, %iterArg_17, %iterArg_18, %iterArg_19, %iterArg_20, %iterArg_21, %iterArg_22, %iterArg_23, %iterArg_24, %iterArg_25, %iterArg_26, %iterArg_27, %iterArg_28, %iterArg_29, %iterArg_30, %iterArg_31, %iterArg_32, %iterArg_33, %iterArg_34, %iterArg_35, %iterArg_36, %iterArg_37, %iterArg_38, %iterArg_39, %iterArg_40, %iterArg_41, %iterArg_42, %iterArg_43, %iterArg_44, %iterArg_45, %iterArg_46, %iterArg_47, %iterArg_48, %iterArg_49, %iterArg_50, %iterArg_51, %iterArg_52, %iterArg_53, %iterArg_54, %iterArg_55, %iterArg_56, %iterArg_57, %iterArg_58, %iterArg_59, %iterArg_60, %iterArg_61, %iterArg_62, %iterArg_63, %iterArg_64, %iterArg_65, %iterArg_66, %iterArg_67, %iterArg_68, %iterArg_69, %iterArg_70, %iterArg_71, %iterArg_72, %iterArg_73, %iterArg_74, %iterArg_75, %iterArg_76, %iterArg_77, %iterArg_78, %iterArg_79, %iterArg_80, %iterArg_81, %iterArg_82, %iterArg_83, %iterArg_84, %iterArg_85, %iterArg_86, %iterArg_87, %iterArg_88, %iterArg_89, %iterArg_90, %iterArg_91, %iterArg_92, %iterArg_93, %iterArg_94, %iterArg_95, %iterArg_96, %iterArg_97, %iterArg_98, %iterArg_99, %iterArg_100, %iterArg_101, %iterArg_102, %iterArg_103, %iterArg_104, %iterArg_105, %iterArg_106, %iterArg_107, %iterArg_108, %iterArg_109, %iterArg_110, %iterArg_111, %iterArg_112, %iterArg_113, %iterArg_114, %iterArg_115, %iterArg_116, %iterArg_117, %iterArg_118, %iterArg_119, %iterArg_120, %iterArg_121, %iterArg_122, %iterArg_123, %iterArg_124, %iterArg_125, %iterArg_126, %iterArg_127, %iterArg_128, %iterArg_129, %iterArg_130, %iterArg_131, %iterArg_132, %iterArg_133, %iterArg_134, %iterArg_135, %iterArg_136, %iterArg_137, %iterArg_138, %iterArg_139, %iterArg_140, %iterArg_141, %iterArg_142, %iterArg_143, %iterArg_144, %iterArg_145, %iterArg_146, %iterArg_147, %iterArg_148, %iterArg_149, %iterArg_150, %iterArg_151, %iterArg_152, %iterArg_153, %iterArg_154, %iterArg_155, %iterArg_156, %iterArg_157, %iterArg_158, %iterArg_159, %iterArg_160, %iterArg_161, %iterArg_162, %iterArg_163, %iterArg_164, %iterArg_165, %iterArg_166, %iterArg_167, %iterArg_168, %iterArg_169, %iterArg_170, %iterArg_171, %iterArg_172, %iterArg_173, %iterArg_174, %iterArg_175, %iterArg_176, %iterArg_177, %iterArg_178, %iterArg_179, %iterArg_180, %iterArg_181, %iterArg_182, %iterArg_183, %iterArg_184, %iterArg_185, %iterArg_186, %iterArg_187, %iterArg_188, %iterArg_189, %iterArg_190, %iterArg_191, %iterArg_192, %iterArg_193, %iterArg_194, %iterArg_195, %iterArg_196, %iterArg_197, %iterArg_198, %iterArg_199, %iterArg_200, %iterArg_201, %iterArg_202, %iterArg_203, %iterArg_204, %iterArg_205, %iterArg_206, %iterArg_207, %iterArg_208, %iterArg_209, %iterArg_210, %iterArg_211, %iterArg_212, %iterArg_213, %iterArg_214, %iterArg_215, %iterArg_216, %iterArg_217, %iterArg_218, %iterArg_219, %iterArg_220, %iterArg_221, %iterArg_222, %iterArg_223, %iterArg_224, %iterArg_225, %iterArg_226, %iterArg_227, %39095, %38917, %38918, %38919, %38920, %38921, %38922, %38923, %38924, %38925, %38926, %38927, %38928, %38929, %38930, %38931, %38932, %38933, %38934, %38935, %38936, %38937, %38938, %38939, %38940, %38941, %38942, %38943, %38944, %38945, %38946, %38947, %38948, %38949, %38950, %38951, %38952, %38953, %38954, %38955, %38956, %38957, %38958, %38959, %38960, %38961, %38962, %38963, %38964, %38965, %38966, %38967, %38968, %38969, %38970, %38971, %38972, %38973, %38974, %38975, %38976, %38977, %38978, %38979, %38980, %38981, %38982, %38983, %38984, %38985, %38986, %38987, %38988, %38989, %38990, %38991, %38992, %38993, %38994, %38995, %38996, %38997, %38998, %38999, %39000, %39001, %39002, %39003, %39004, %39005, %39006, %39007, %39008, %39009, %39010, %39011, %39012, %39013, %39014, %39015, %39016, %39017, %39018, %39019, %39020, %39021, %39022, %39023, %39024, %39025, %39026, %39027, %39028, %39029, %39030, %39031, %39032, %39033, %39034, %39035, %39036, %39037, %39038, %39039, %39040, %39041, %39042, %39043, %39044, %39045, %39046, %39047, %39048, %39049, %39050, %39051, %39052, %39053, %39054, %39055, %39056, %39057, %39058, %39059, %39060, %39061, %39062, %39063, %39064, %39065, %39066, %39067, %39068, %39069, %39070, %39071, %39072, %39073, %39074, %39075, %39076, %39077, %39078, %39079, %39080, %39081, %39082, %39083, %39084, %39085, %39086, %39087, %39088, %39089, %39090, %39091, %39092, %39093, %38913#1 : tensor<1x1x1x2048x2048xbf16>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1x1x2048x2048xbf16>, tensor<1xi32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x1x2048xi32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<i32>, tensor<1x1x1x2048x2048xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x8192xbf16>, tensor<1x2048x8192xbf16>, tensor<1x2048xbf16>, tensor<1x8192x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x2048xbf16>, tensor<1x3x2048x32x64xbf16>, tensor<1x64xbf16>, tensor<1x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16> loc(#loc1971)
    } loc(#loc1971)
    %1707 = stablehlo.convert %1706#406 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1708 = stablehlo.convert %1706#405 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1709 = stablehlo.convert %1706#404 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1710 = stablehlo.convert %1706#403 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1711 = stablehlo.convert %1706#402 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1712 = stablehlo.convert %1706#401 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1713 = stablehlo.convert %1706#400 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1714 = stablehlo.convert %1706#399 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1715 = stablehlo.convert %1706#398 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1716 = stablehlo.convert %1706#397 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1717 = stablehlo.convert %1706#396 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1718 = stablehlo.convert %1706#395 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1719 = stablehlo.convert %1706#394 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1720 = stablehlo.convert %1706#393 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1721 = stablehlo.convert %1706#392 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1722 = stablehlo.convert %1706#391 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1723 = stablehlo.convert %1706#390 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1724 = stablehlo.convert %1706#389 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1725 = stablehlo.convert %1706#388 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1726 = stablehlo.convert %1706#387 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1727 = stablehlo.convert %1706#386 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1728 = stablehlo.convert %1706#385 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1729 = stablehlo.convert %1706#384 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1730 = stablehlo.convert %1706#383 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1731 = stablehlo.convert %1706#382 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1732 = stablehlo.convert %1706#381 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1733 = stablehlo.convert %1706#380 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1734 = stablehlo.convert %1706#379 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1735 = stablehlo.convert %1706#378 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1736 = stablehlo.convert %1706#377 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1737 = stablehlo.convert %1706#376 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1738 = stablehlo.convert %1706#375 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1739 = stablehlo.convert %1706#374 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1740 = stablehlo.convert %1706#373 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1741 = stablehlo.convert %1706#372 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1742 = stablehlo.convert %1706#371 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1743 = stablehlo.convert %1706#370 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1744 = stablehlo.convert %1706#369 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1745 = stablehlo.convert %1706#368 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1746 = stablehlo.convert %1706#367 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1747 = stablehlo.convert %1706#366 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1748 = stablehlo.convert %1706#365 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1749 = stablehlo.convert %1706#364 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1750 = stablehlo.convert %1706#363 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1751 = stablehlo.convert %1706#362 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1752 = stablehlo.convert %1706#361 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1753 = stablehlo.convert %1706#360 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1754 = stablehlo.convert %1706#359 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1755 = stablehlo.convert %1706#358 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1756 = stablehlo.convert %1706#357 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1757 = stablehlo.convert %1706#356 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1758 = stablehlo.convert %1706#355 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1759 = stablehlo.convert %1706#354 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1760 = stablehlo.convert %1706#353 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1761 = stablehlo.convert %1706#352 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1762 = stablehlo.convert %1706#351 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1763 = stablehlo.convert %1706#350 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1764 = stablehlo.convert %1706#349 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1765 = stablehlo.convert %1706#348 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1766 = stablehlo.convert %1706#347 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1767 = stablehlo.convert %1706#346 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1768 = stablehlo.convert %1706#345 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1769 = stablehlo.convert %1706#344 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1770 = stablehlo.convert %1706#343 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1771 = stablehlo.convert %1706#342 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1772 = stablehlo.convert %1706#341 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1773 = stablehlo.convert %1706#340 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1774 = stablehlo.convert %1706#339 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1775 = stablehlo.convert %1706#338 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1776 = stablehlo.convert %1706#337 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1777 = stablehlo.convert %1706#336 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1778 = stablehlo.convert %1706#335 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1779 = stablehlo.convert %1706#334 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1780 = stablehlo.convert %1706#333 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1781 = stablehlo.convert %1706#332 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1782 = stablehlo.convert %1706#331 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1783 = stablehlo.convert %1706#330 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1784 = stablehlo.convert %1706#329 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1785 = stablehlo.convert %1706#328 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1786 = stablehlo.convert %1706#327 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1787 = stablehlo.convert %1706#326 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1788 = stablehlo.convert %1706#325 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1789 = stablehlo.convert %1706#324 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1790 = stablehlo.convert %1706#323 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1791 = stablehlo.convert %1706#322 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1792 = stablehlo.convert %1706#321 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1793 = stablehlo.convert %1706#320 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1794 = stablehlo.convert %1706#319 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1795 = stablehlo.convert %1706#318 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1796 = stablehlo.convert %1706#317 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1797 = stablehlo.convert %1706#316 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1798 = stablehlo.convert %1706#315 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1799 = stablehlo.convert %1706#314 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1800 = stablehlo.convert %1706#313 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1801 = stablehlo.convert %1706#312 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1802 = stablehlo.convert %1706#311 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1803 = stablehlo.convert %1706#310 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1804 = stablehlo.convert %1706#309 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1805 = stablehlo.convert %1706#308 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1806 = stablehlo.convert %1706#307 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1807 = stablehlo.convert %1706#306 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1808 = stablehlo.convert %1706#305 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1809 = stablehlo.convert %1706#304 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1810 = stablehlo.convert %1706#303 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1811 = stablehlo.convert %1706#302 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1812 = stablehlo.convert %1706#301 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1813 = stablehlo.convert %1706#300 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1814 = stablehlo.convert %1706#299 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1815 = stablehlo.convert %1706#298 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1816 = stablehlo.convert %1706#297 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1817 = stablehlo.convert %1706#296 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1818 = stablehlo.convert %1706#295 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1819 = stablehlo.convert %1706#294 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1820 = stablehlo.convert %1706#293 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1821 = stablehlo.convert %1706#292 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1822 = stablehlo.convert %1706#291 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1823 = stablehlo.convert %1706#290 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1824 = stablehlo.convert %1706#289 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1825 = stablehlo.convert %1706#288 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1826 = stablehlo.convert %1706#287 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1827 = stablehlo.convert %1706#286 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1828 = stablehlo.convert %1706#285 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1829 = stablehlo.convert %1706#284 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1830 = stablehlo.convert %1706#283 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1831 = stablehlo.convert %1706#282 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1832 = stablehlo.convert %1706#281 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1833 = stablehlo.convert %1706#280 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1834 = stablehlo.convert %1706#279 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1835 = stablehlo.convert %1706#278 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1836 = stablehlo.convert %1706#277 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1837 = stablehlo.convert %1706#276 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1838 = stablehlo.convert %1706#275 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1839 = stablehlo.convert %1706#274 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1840 = stablehlo.convert %1706#273 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1841 = stablehlo.convert %1706#272 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1842 = stablehlo.convert %1706#271 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1843 = stablehlo.convert %1706#270 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1844 = stablehlo.convert %1706#269 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1845 = stablehlo.convert %1706#268 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1846 = stablehlo.convert %1706#267 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1847 = stablehlo.convert %1706#266 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1848 = stablehlo.convert %1706#265 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1849 = stablehlo.convert %1706#264 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1850 = stablehlo.convert %1706#263 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1851 = stablehlo.convert %1706#262 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1852 = stablehlo.convert %1706#261 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1853 = stablehlo.convert %1706#260 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1854 = stablehlo.convert %1706#259 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1855 = stablehlo.convert %1706#258 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1856 = stablehlo.convert %1706#257 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1857 = stablehlo.convert %1706#256 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1858 = stablehlo.convert %1706#255 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1859 = stablehlo.convert %1706#254 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1860 = stablehlo.convert %1706#253 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1861 = stablehlo.convert %1706#252 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1862 = stablehlo.convert %1706#251 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1863 = stablehlo.convert %1706#250 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1864 = stablehlo.convert %1706#249 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1865 = stablehlo.convert %1706#248 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1866 = stablehlo.convert %1706#247 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1867 = stablehlo.convert %1706#246 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1868 = stablehlo.convert %1706#245 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1869 = stablehlo.convert %1706#244 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1870 = stablehlo.convert %1706#243 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1871 = stablehlo.convert %1706#242 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1872 = stablehlo.convert %1706#241 : (tensor<1x2048x32x64xbf16>) -> tensor<1x2048x32x64xf32> loc(#loc1929)
    %1873 = stablehlo.convert %1706#240 : (tensor<1x64xbf16>) -> tensor<1x64xf32> loc(#loc1929)
    %1874 = stablehlo.convert %1706#239 : (tensor<1x3x2048x32x64xbf16>) -> tensor<1x3x2048x32x64xf32> loc(#loc1929)
    %1875 = stablehlo.convert %1706#238 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1876 = stablehlo.convert %1706#237 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1877 = stablehlo.convert %1706#236 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1878 = stablehlo.convert %1706#235 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1879 = stablehlo.convert %1706#234 : (tensor<1x8192x2048xbf16>) -> tensor<1x8192x2048xf32> loc(#loc1929)
    %1880 = stablehlo.convert %1706#233 : (tensor<1x2048xbf16>) -> tensor<1x2048xf32> loc(#loc1929)
    %1881 = stablehlo.convert %1706#232 : (tensor<1x2048x8192xbf16>) -> tensor<1x2048x8192xf32> loc(#loc1929)
    %1882 = stablehlo.convert %1706#231 : (tensor<1x8192xbf16>) -> tensor<1x8192xf32> loc(#loc1929)
    %1883 = call @_pad_1137(%1706#230) : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5211)
    %1884 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5212)
    %1885 = stablehlo.reduce(%1883 init: %1884) across dimensions = [0, 1, 2] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2048x2048xbf16>
     reducer(%arg558: tensor<bf16> loc(unknown), %arg559: tensor<bf16> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<bf16> loc(#loc5212)
      stablehlo.return %29929 : tensor<bf16> loc(#loc5212)
    } loc(#loc5212)
    %1886 = stablehlo.broadcast_in_dim %1885, dims = [2, 3] : (tensor<2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5213)
    %1887 = stablehlo.reshape %1886 : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5214)
    %1888 = stablehlo.transpose %1887, dims = [0, 2, 1, 3, 4] : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5215)
    %1889 = stablehlo.reshape %1888 : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc5216)
    %1890 = arith.select %ptrue, %1889, %1889 : tensor<1x2048x2048xbf16>
    %1891 = arith.select %ptrue, %1890, %1890 : tensor<1x2048x2048xbf16>
    %1892 = call @_einsum_1138(%209, %1891) : (tensor<1x2048x2048xbf16>, tensor<1x2048x2048xbf16>) -> tensor<2048x2048xbf16> loc(#loc5218)
    %1893 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %1894 = stablehlo.pad %1892, %1893, low = [0, 0], high = [0, 0], interior = [0, 0] : (tensor<2048x2048xbf16>, tensor<bf16>) -> tensor<2048x2048xbf16> loc(#loc5219)
    %1895 = stablehlo.convert %1894 : (tensor<2048x2048xbf16>) -> tensor<2048x2048xf32> loc(#loc1929)
    %1896 = arith.select %ptrue, %1890, %1890 : tensor<1x2048x2048xbf16>
    %1897 = stablehlo.constant dense<4.525000e+01> : tensor<bf16> loc(#loc)
    %1898 = stablehlo.broadcast_in_dim %1897, dims = [] : (tensor<bf16>) -> tensor<1x2048x2048xbf16> loc(#loc5221)
    %1899 = stablehlo.multiply %1896, %1898 : tensor<1x2048x2048xbf16> loc(#loc5221)
    %1900 = call @_einsum_1139(%202, %1899) : (tensor<1x2048x51200xbf16>, tensor<1x2048x2048xbf16>) -> tensor<51200x2048xbf16> loc(#loc5222)
    %1901 = stablehlo.transpose %1900, dims = [1, 0] : (tensor<51200x2048xbf16>) -> tensor<2048x51200xbf16> loc(#loc5223)
    %1902 = stablehlo.add %1192#1, %1901 : tensor<2048x51200xbf16> loc(#loc5224)
    %1903 = stablehlo.convert %1902 : (tensor<2048x51200xbf16>) -> tensor<2048x51200xf32> loc(#loc1929)
    %1904 = stablehlo.convert %arg0 : (tensor<ui32>) -> tensor<f32> loc(#loc5225)
    %1905 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1906 = stablehlo.subtract %1904, %1905 : tensor<f32> loc(#loc5226)
    %1907 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1908 = stablehlo.divide %1906, %1907 : tensor<f32> loc(#loc5227)
    %1909 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1910 = stablehlo.multiply %1908, %1909 : tensor<f32> loc(#loc5228)
    %1911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1912 = stablehlo.add %1911, %1910 : tensor<f32> loc(#loc5229)
    %1913 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1914 = stablehlo.compare  LT, %1904, %1913,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5230)
    %1915 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1916 = stablehlo.compare  GE, %1904, %1915,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5231)
    %1917 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1918 = call @_where_1140(%1916, %1917, %1912) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %1919 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1920 = call @_where_1141(%1914, %1919, %1918) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %1921 = stablehlo.constant dense<1> : tensor<ui32> loc(#loc)
    %1922 = stablehlo.compare  LT, %arg0, %1921,  UNSIGNED : (tensor<ui32>, tensor<ui32>) -> tensor<i1> loc(#loc5233)
    %1923 = stablehlo.constant dense<1> : tensor<ui32> loc(#loc)
    %1924 = stablehlo.subtract %arg0, %1923 : tensor<ui32> loc(#loc5234)
    %1925 = stablehlo.convert %1924 : (tensor<ui32>) -> tensor<f32> loc(#loc5225)
    %1926 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1927 = stablehlo.subtract %1925, %1926 : tensor<f32> loc(#loc5226)
    %1928 = stablehlo.constant dense<4.999990e+05> : tensor<f32> loc(#loc)
    %1929 = stablehlo.divide %1927, %1928 : tensor<f32> loc(#loc5227)
    %1930 = stablehlo.constant dense<3.14159274> : tensor<f32> loc(#loc)
    %1931 = stablehlo.multiply %1929, %1930 : tensor<f32> loc(#loc5228)
    %1932 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1933 = stablehlo.add %1932, %1931 : tensor<f32> loc(#loc5229)
    %1934 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %1935 = stablehlo.compare  LT, %1925, %1934,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5230)
    %1936 = stablehlo.constant dense<4.999990e+05> : tensor<f32> loc(#loc)
    %1937 = stablehlo.compare  GE, %1925, %1936,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5231)
    %1938 = stablehlo.constant dense<3.14159274> : tensor<f32> loc(#loc)
    %1939 = call @_where_1142(%1937, %1938, %1933) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %1940 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %1941 = call @_where_1143(%1935, %1940, %1939) : (tensor<i1>, tensor<i32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %1942 = stablehlo.cosine %1941 : tensor<f32> loc(#loc5235)
    %1943 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %1944 = stablehlo.add %1943, %1942 : tensor<f32> loc(#loc5229)
    %1945 = stablehlo.constant dense<4.500000e-01> : tensor<f32> loc(#loc)
    %1946 = stablehlo.multiply %1945, %1944 : tensor<f32> loc(#loc5236)
    %1947 = stablehlo.constant dense<1.000000e-01> : tensor<f32> loc(#loc)
    %1948 = stablehlo.add %1947, %1946 : tensor<f32> loc(#loc5229)
    %1949 = call @_where_1144(%1922, %1920, %1948) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5237)
    %1950 = stablehlo.constant dense<6.000000e-04> : tensor<f32> loc(#loc)
    %1951 = stablehlo.multiply %1949, %1950 : tensor<f32> loc(#loc5228)
    %1952 = stablehlo.multiply %1198, %1198 : tensor<2048xf32> loc(#loc5238)
    %1953 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %1954 = stablehlo.reduce(%1952 init: %1953) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %1955 = stablehlo.multiply %1207, %1207 : tensor<2048xf32> loc(#loc5238)
    %1956 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %1957 = stablehlo.reduce(%1955 init: %1956) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %1958 = stablehlo.multiply %1895, %1895 : tensor<2048x2048xf32> loc(#loc5238)
    %1959 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %1960 = stablehlo.reduce(%1958 init: %1959) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %1961 = stablehlo.multiply %1190, %1190 : tensor<51200xf32> loc(#loc5238)
    %1962 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %1963 = stablehlo.reduce(%1961 init: %1962) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %1964 = stablehlo.multiply %1903, %1903 : tensor<2048x51200xf32> loc(#loc5238)
    %1965 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %1966 = stablehlo.reduce(%1964 init: %1965) across dimensions = [0, 1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %1967 = stablehlo.multiply %1882, %1882 : tensor<1x8192xf32> loc(#loc5238)
    %1968 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %1969 = stablehlo.reduce(%1967 init: %1968) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %1970 = stablehlo.multiply %1881, %1881 : tensor<1x2048x8192xf32> loc(#loc5238)
    %1971 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %1972 = stablehlo.reduce(%1970 init: %1971) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %1973 = stablehlo.multiply %1880, %1880 : tensor<1x2048xf32> loc(#loc5242)
    %1974 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %1975 = stablehlo.reduce(%1973 init: %1974) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %1976 = stablehlo.multiply %1879, %1879 : tensor<1x8192x2048xf32> loc(#loc5238)
    %1977 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %1978 = stablehlo.reduce(%1976 init: %1977) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %1979 = stablehlo.multiply %1878, %1878 : tensor<1x2048xf32> loc(#loc5242)
    %1980 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %1981 = stablehlo.reduce(%1979 init: %1980) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %1982 = stablehlo.multiply %1877, %1877 : tensor<1x2048xf32> loc(#loc5242)
    %1983 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %1984 = stablehlo.reduce(%1982 init: %1983) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %1985 = stablehlo.multiply %1876, %1876 : tensor<1x2048xf32> loc(#loc5242)
    %1986 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %1987 = stablehlo.reduce(%1985 init: %1986) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %1988 = stablehlo.multiply %1875, %1875 : tensor<1x2048xf32> loc(#loc5242)
    %1989 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %1990 = stablehlo.reduce(%1988 init: %1989) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %1991 = stablehlo.multiply %1874, %1874 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %1992 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %1993 = stablehlo.reduce(%1991 init: %1992) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %1994 = stablehlo.multiply %1873, %1873 : tensor<1x64xf32> loc(#loc5238)
    %1995 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %1996 = stablehlo.reduce(%1994 init: %1995) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %1997 = stablehlo.multiply %1872, %1872 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %1998 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %1999 = stablehlo.reduce(%1997 init: %1998) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2000 = stablehlo.multiply %1871, %1871 : tensor<1x8192xf32> loc(#loc5238)
    %2001 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2002 = stablehlo.reduce(%2000 init: %2001) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2003 = stablehlo.multiply %1870, %1870 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2004 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2005 = stablehlo.reduce(%2003 init: %2004) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2006 = stablehlo.multiply %1869, %1869 : tensor<1x2048xf32> loc(#loc5242)
    %2007 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2008 = stablehlo.reduce(%2006 init: %2007) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2009 = stablehlo.multiply %1868, %1868 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2010 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2011 = stablehlo.reduce(%2009 init: %2010) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2012 = stablehlo.multiply %1867, %1867 : tensor<1x2048xf32> loc(#loc5242)
    %2013 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2014 = stablehlo.reduce(%2012 init: %2013) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2015 = stablehlo.multiply %1866, %1866 : tensor<1x2048xf32> loc(#loc5242)
    %2016 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2017 = stablehlo.reduce(%2015 init: %2016) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2018 = stablehlo.multiply %1865, %1865 : tensor<1x2048xf32> loc(#loc5242)
    %2019 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2020 = stablehlo.reduce(%2018 init: %2019) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2021 = stablehlo.multiply %1864, %1864 : tensor<1x2048xf32> loc(#loc5242)
    %2022 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2023 = stablehlo.reduce(%2021 init: %2022) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2024 = stablehlo.multiply %1863, %1863 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2025 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2026 = stablehlo.reduce(%2024 init: %2025) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2027 = stablehlo.multiply %1862, %1862 : tensor<1x64xf32> loc(#loc5238)
    %2028 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2029 = stablehlo.reduce(%2027 init: %2028) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2030 = stablehlo.multiply %1861, %1861 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2031 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2032 = stablehlo.reduce(%2030 init: %2031) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2033 = stablehlo.multiply %1860, %1860 : tensor<1x8192xf32> loc(#loc5238)
    %2034 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2035 = stablehlo.reduce(%2033 init: %2034) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2036 = stablehlo.multiply %1859, %1859 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2037 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2038 = stablehlo.reduce(%2036 init: %2037) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2039 = stablehlo.multiply %1858, %1858 : tensor<1x2048xf32> loc(#loc5242)
    %2040 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2041 = stablehlo.reduce(%2039 init: %2040) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2042 = stablehlo.multiply %1857, %1857 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2043 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2044 = stablehlo.reduce(%2042 init: %2043) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2045 = stablehlo.multiply %1856, %1856 : tensor<1x2048xf32> loc(#loc5242)
    %2046 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2047 = stablehlo.reduce(%2045 init: %2046) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2048 = stablehlo.multiply %1855, %1855 : tensor<1x2048xf32> loc(#loc5242)
    %2049 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2050 = stablehlo.reduce(%2048 init: %2049) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2051 = stablehlo.multiply %1854, %1854 : tensor<1x2048xf32> loc(#loc5242)
    %2052 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2053 = stablehlo.reduce(%2051 init: %2052) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2054 = stablehlo.multiply %1853, %1853 : tensor<1x2048xf32> loc(#loc5242)
    %2055 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2056 = stablehlo.reduce(%2054 init: %2055) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2057 = stablehlo.multiply %1852, %1852 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2058 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2059 = stablehlo.reduce(%2057 init: %2058) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2060 = stablehlo.multiply %1851, %1851 : tensor<1x64xf32> loc(#loc5238)
    %2061 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2062 = stablehlo.reduce(%2060 init: %2061) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2063 = stablehlo.multiply %1850, %1850 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2064 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2065 = stablehlo.reduce(%2063 init: %2064) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2066 = stablehlo.multiply %1849, %1849 : tensor<1x8192xf32> loc(#loc5238)
    %2067 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2068 = stablehlo.reduce(%2066 init: %2067) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2069 = stablehlo.multiply %1848, %1848 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2070 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2071 = stablehlo.reduce(%2069 init: %2070) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2072 = stablehlo.multiply %1847, %1847 : tensor<1x2048xf32> loc(#loc5242)
    %2073 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2074 = stablehlo.reduce(%2072 init: %2073) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2075 = stablehlo.multiply %1846, %1846 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2076 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2077 = stablehlo.reduce(%2075 init: %2076) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2078 = stablehlo.multiply %1845, %1845 : tensor<1x2048xf32> loc(#loc5242)
    %2079 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2080 = stablehlo.reduce(%2078 init: %2079) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2081 = stablehlo.multiply %1844, %1844 : tensor<1x2048xf32> loc(#loc5242)
    %2082 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2083 = stablehlo.reduce(%2081 init: %2082) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2084 = stablehlo.multiply %1843, %1843 : tensor<1x2048xf32> loc(#loc5242)
    %2085 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2086 = stablehlo.reduce(%2084 init: %2085) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2087 = stablehlo.multiply %1842, %1842 : tensor<1x2048xf32> loc(#loc5242)
    %2088 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2089 = stablehlo.reduce(%2087 init: %2088) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2090 = stablehlo.multiply %1841, %1841 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2091 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2092 = stablehlo.reduce(%2090 init: %2091) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2093 = stablehlo.multiply %1840, %1840 : tensor<1x64xf32> loc(#loc5238)
    %2094 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2095 = stablehlo.reduce(%2093 init: %2094) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2096 = stablehlo.multiply %1839, %1839 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2097 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2098 = stablehlo.reduce(%2096 init: %2097) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2099 = stablehlo.multiply %1838, %1838 : tensor<1x8192xf32> loc(#loc5238)
    %2100 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2101 = stablehlo.reduce(%2099 init: %2100) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2102 = stablehlo.multiply %1837, %1837 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2103 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2104 = stablehlo.reduce(%2102 init: %2103) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2105 = stablehlo.multiply %1836, %1836 : tensor<1x2048xf32> loc(#loc5242)
    %2106 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2107 = stablehlo.reduce(%2105 init: %2106) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2108 = stablehlo.multiply %1835, %1835 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2109 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2110 = stablehlo.reduce(%2108 init: %2109) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2111 = stablehlo.multiply %1834, %1834 : tensor<1x2048xf32> loc(#loc5242)
    %2112 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2113 = stablehlo.reduce(%2111 init: %2112) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2114 = stablehlo.multiply %1833, %1833 : tensor<1x2048xf32> loc(#loc5242)
    %2115 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2116 = stablehlo.reduce(%2114 init: %2115) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2117 = stablehlo.multiply %1832, %1832 : tensor<1x2048xf32> loc(#loc5242)
    %2118 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2119 = stablehlo.reduce(%2117 init: %2118) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2120 = stablehlo.multiply %1831, %1831 : tensor<1x2048xf32> loc(#loc5242)
    %2121 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2122 = stablehlo.reduce(%2120 init: %2121) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2123 = stablehlo.multiply %1830, %1830 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2124 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2125 = stablehlo.reduce(%2123 init: %2124) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2126 = stablehlo.multiply %1829, %1829 : tensor<1x64xf32> loc(#loc5238)
    %2127 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2128 = stablehlo.reduce(%2126 init: %2127) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2129 = stablehlo.multiply %1828, %1828 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2130 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2131 = stablehlo.reduce(%2129 init: %2130) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2132 = stablehlo.multiply %1827, %1827 : tensor<1x8192xf32> loc(#loc5238)
    %2133 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2134 = stablehlo.reduce(%2132 init: %2133) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2135 = stablehlo.multiply %1826, %1826 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2136 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2137 = stablehlo.reduce(%2135 init: %2136) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2138 = stablehlo.multiply %1825, %1825 : tensor<1x2048xf32> loc(#loc5242)
    %2139 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2140 = stablehlo.reduce(%2138 init: %2139) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2141 = stablehlo.multiply %1824, %1824 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2142 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2143 = stablehlo.reduce(%2141 init: %2142) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2144 = stablehlo.multiply %1823, %1823 : tensor<1x2048xf32> loc(#loc5242)
    %2145 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2146 = stablehlo.reduce(%2144 init: %2145) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2147 = stablehlo.multiply %1822, %1822 : tensor<1x2048xf32> loc(#loc5242)
    %2148 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2149 = stablehlo.reduce(%2147 init: %2148) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2150 = stablehlo.multiply %1821, %1821 : tensor<1x2048xf32> loc(#loc5242)
    %2151 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2152 = stablehlo.reduce(%2150 init: %2151) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2153 = stablehlo.multiply %1820, %1820 : tensor<1x2048xf32> loc(#loc5242)
    %2154 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2155 = stablehlo.reduce(%2153 init: %2154) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2156 = stablehlo.multiply %1819, %1819 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2157 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2158 = stablehlo.reduce(%2156 init: %2157) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2159 = stablehlo.multiply %1818, %1818 : tensor<1x64xf32> loc(#loc5238)
    %2160 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2161 = stablehlo.reduce(%2159 init: %2160) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2162 = stablehlo.multiply %1817, %1817 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2163 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2164 = stablehlo.reduce(%2162 init: %2163) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2165 = stablehlo.multiply %1816, %1816 : tensor<1x8192xf32> loc(#loc5238)
    %2166 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2167 = stablehlo.reduce(%2165 init: %2166) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2168 = stablehlo.multiply %1815, %1815 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2169 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2170 = stablehlo.reduce(%2168 init: %2169) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2171 = stablehlo.multiply %1814, %1814 : tensor<1x2048xf32> loc(#loc5242)
    %2172 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2173 = stablehlo.reduce(%2171 init: %2172) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2174 = stablehlo.multiply %1813, %1813 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2175 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2176 = stablehlo.reduce(%2174 init: %2175) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2177 = stablehlo.multiply %1812, %1812 : tensor<1x2048xf32> loc(#loc5242)
    %2178 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2179 = stablehlo.reduce(%2177 init: %2178) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2180 = stablehlo.multiply %1811, %1811 : tensor<1x2048xf32> loc(#loc5242)
    %2181 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2182 = stablehlo.reduce(%2180 init: %2181) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2183 = stablehlo.multiply %1810, %1810 : tensor<1x2048xf32> loc(#loc5242)
    %2184 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2185 = stablehlo.reduce(%2183 init: %2184) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2186 = stablehlo.multiply %1809, %1809 : tensor<1x2048xf32> loc(#loc5242)
    %2187 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2188 = stablehlo.reduce(%2186 init: %2187) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2189 = stablehlo.multiply %1808, %1808 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2190 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2191 = stablehlo.reduce(%2189 init: %2190) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2192 = stablehlo.multiply %1807, %1807 : tensor<1x64xf32> loc(#loc5238)
    %2193 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2194 = stablehlo.reduce(%2192 init: %2193) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2195 = stablehlo.multiply %1806, %1806 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2196 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2197 = stablehlo.reduce(%2195 init: %2196) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2198 = stablehlo.multiply %1805, %1805 : tensor<1x8192xf32> loc(#loc5238)
    %2199 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2200 = stablehlo.reduce(%2198 init: %2199) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2201 = stablehlo.multiply %1804, %1804 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2202 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2203 = stablehlo.reduce(%2201 init: %2202) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2204 = stablehlo.multiply %1803, %1803 : tensor<1x2048xf32> loc(#loc5242)
    %2205 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2206 = stablehlo.reduce(%2204 init: %2205) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2207 = stablehlo.multiply %1802, %1802 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2208 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2209 = stablehlo.reduce(%2207 init: %2208) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2210 = stablehlo.multiply %1801, %1801 : tensor<1x2048xf32> loc(#loc5242)
    %2211 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2212 = stablehlo.reduce(%2210 init: %2211) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2213 = stablehlo.multiply %1800, %1800 : tensor<1x2048xf32> loc(#loc5242)
    %2214 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2215 = stablehlo.reduce(%2213 init: %2214) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2216 = stablehlo.multiply %1799, %1799 : tensor<1x2048xf32> loc(#loc5242)
    %2217 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2218 = stablehlo.reduce(%2216 init: %2217) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2219 = stablehlo.multiply %1798, %1798 : tensor<1x2048xf32> loc(#loc5242)
    %2220 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2221 = stablehlo.reduce(%2219 init: %2220) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2222 = stablehlo.multiply %1797, %1797 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2223 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2224 = stablehlo.reduce(%2222 init: %2223) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2225 = stablehlo.multiply %1796, %1796 : tensor<1x64xf32> loc(#loc5238)
    %2226 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2227 = stablehlo.reduce(%2225 init: %2226) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2228 = stablehlo.multiply %1795, %1795 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2229 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2230 = stablehlo.reduce(%2228 init: %2229) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2231 = stablehlo.multiply %1794, %1794 : tensor<1x8192xf32> loc(#loc5238)
    %2232 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2233 = stablehlo.reduce(%2231 init: %2232) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2234 = stablehlo.multiply %1793, %1793 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2235 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2236 = stablehlo.reduce(%2234 init: %2235) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2237 = stablehlo.multiply %1792, %1792 : tensor<1x2048xf32> loc(#loc5242)
    %2238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2239 = stablehlo.reduce(%2237 init: %2238) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2240 = stablehlo.multiply %1791, %1791 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2241 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2242 = stablehlo.reduce(%2240 init: %2241) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2243 = stablehlo.multiply %1790, %1790 : tensor<1x2048xf32> loc(#loc5242)
    %2244 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2245 = stablehlo.reduce(%2243 init: %2244) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2246 = stablehlo.multiply %1789, %1789 : tensor<1x2048xf32> loc(#loc5242)
    %2247 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2248 = stablehlo.reduce(%2246 init: %2247) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2249 = stablehlo.multiply %1788, %1788 : tensor<1x2048xf32> loc(#loc5242)
    %2250 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2251 = stablehlo.reduce(%2249 init: %2250) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2252 = stablehlo.multiply %1787, %1787 : tensor<1x2048xf32> loc(#loc5242)
    %2253 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2254 = stablehlo.reduce(%2252 init: %2253) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2255 = stablehlo.multiply %1786, %1786 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2256 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2257 = stablehlo.reduce(%2255 init: %2256) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2258 = stablehlo.multiply %1785, %1785 : tensor<1x64xf32> loc(#loc5238)
    %2259 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2260 = stablehlo.reduce(%2258 init: %2259) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2261 = stablehlo.multiply %1784, %1784 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2262 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2263 = stablehlo.reduce(%2261 init: %2262) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2264 = stablehlo.multiply %1783, %1783 : tensor<1x8192xf32> loc(#loc5238)
    %2265 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2266 = stablehlo.reduce(%2264 init: %2265) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2267 = stablehlo.multiply %1782, %1782 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2268 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2269 = stablehlo.reduce(%2267 init: %2268) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2270 = stablehlo.multiply %1781, %1781 : tensor<1x2048xf32> loc(#loc5242)
    %2271 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2272 = stablehlo.reduce(%2270 init: %2271) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2273 = stablehlo.multiply %1780, %1780 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2274 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2275 = stablehlo.reduce(%2273 init: %2274) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2276 = stablehlo.multiply %1779, %1779 : tensor<1x2048xf32> loc(#loc5242)
    %2277 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2278 = stablehlo.reduce(%2276 init: %2277) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2279 = stablehlo.multiply %1778, %1778 : tensor<1x2048xf32> loc(#loc5242)
    %2280 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2281 = stablehlo.reduce(%2279 init: %2280) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2282 = stablehlo.multiply %1777, %1777 : tensor<1x2048xf32> loc(#loc5242)
    %2283 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2284 = stablehlo.reduce(%2282 init: %2283) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2285 = stablehlo.multiply %1776, %1776 : tensor<1x2048xf32> loc(#loc5242)
    %2286 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2287 = stablehlo.reduce(%2285 init: %2286) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2288 = stablehlo.multiply %1775, %1775 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2289 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2290 = stablehlo.reduce(%2288 init: %2289) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2291 = stablehlo.multiply %1774, %1774 : tensor<1x64xf32> loc(#loc5238)
    %2292 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2293 = stablehlo.reduce(%2291 init: %2292) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2294 = stablehlo.multiply %1773, %1773 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2295 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2296 = stablehlo.reduce(%2294 init: %2295) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2297 = stablehlo.multiply %1772, %1772 : tensor<1x8192xf32> loc(#loc5238)
    %2298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2299 = stablehlo.reduce(%2297 init: %2298) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2300 = stablehlo.multiply %1771, %1771 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2301 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2302 = stablehlo.reduce(%2300 init: %2301) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2303 = stablehlo.multiply %1770, %1770 : tensor<1x2048xf32> loc(#loc5242)
    %2304 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2305 = stablehlo.reduce(%2303 init: %2304) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2306 = stablehlo.multiply %1769, %1769 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2307 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2308 = stablehlo.reduce(%2306 init: %2307) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2309 = stablehlo.multiply %1768, %1768 : tensor<1x2048xf32> loc(#loc5242)
    %2310 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2311 = stablehlo.reduce(%2309 init: %2310) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2312 = stablehlo.multiply %1767, %1767 : tensor<1x2048xf32> loc(#loc5242)
    %2313 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2314 = stablehlo.reduce(%2312 init: %2313) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2315 = stablehlo.multiply %1766, %1766 : tensor<1x2048xf32> loc(#loc5242)
    %2316 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2317 = stablehlo.reduce(%2315 init: %2316) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2318 = stablehlo.multiply %1765, %1765 : tensor<1x2048xf32> loc(#loc5242)
    %2319 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2320 = stablehlo.reduce(%2318 init: %2319) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2321 = stablehlo.multiply %1764, %1764 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2322 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2323 = stablehlo.reduce(%2321 init: %2322) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2324 = stablehlo.multiply %1763, %1763 : tensor<1x64xf32> loc(#loc5238)
    %2325 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2326 = stablehlo.reduce(%2324 init: %2325) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2327 = stablehlo.multiply %1762, %1762 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2328 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2329 = stablehlo.reduce(%2327 init: %2328) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2330 = stablehlo.multiply %1761, %1761 : tensor<1x8192xf32> loc(#loc5238)
    %2331 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2332 = stablehlo.reduce(%2330 init: %2331) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2333 = stablehlo.multiply %1760, %1760 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2334 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2335 = stablehlo.reduce(%2333 init: %2334) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2336 = stablehlo.multiply %1759, %1759 : tensor<1x2048xf32> loc(#loc5242)
    %2337 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2338 = stablehlo.reduce(%2336 init: %2337) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2339 = stablehlo.multiply %1758, %1758 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2340 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2341 = stablehlo.reduce(%2339 init: %2340) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2342 = stablehlo.multiply %1757, %1757 : tensor<1x2048xf32> loc(#loc5242)
    %2343 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2344 = stablehlo.reduce(%2342 init: %2343) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2345 = stablehlo.multiply %1756, %1756 : tensor<1x2048xf32> loc(#loc5242)
    %2346 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2347 = stablehlo.reduce(%2345 init: %2346) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2348 = stablehlo.multiply %1755, %1755 : tensor<1x2048xf32> loc(#loc5242)
    %2349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2350 = stablehlo.reduce(%2348 init: %2349) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2351 = stablehlo.multiply %1754, %1754 : tensor<1x2048xf32> loc(#loc5242)
    %2352 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2353 = stablehlo.reduce(%2351 init: %2352) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2354 = stablehlo.multiply %1753, %1753 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2355 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2356 = stablehlo.reduce(%2354 init: %2355) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2357 = stablehlo.multiply %1752, %1752 : tensor<1x64xf32> loc(#loc5238)
    %2358 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2359 = stablehlo.reduce(%2357 init: %2358) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2360 = stablehlo.multiply %1751, %1751 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2361 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2362 = stablehlo.reduce(%2360 init: %2361) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2363 = stablehlo.multiply %1750, %1750 : tensor<1x8192xf32> loc(#loc5238)
    %2364 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2365 = stablehlo.reduce(%2363 init: %2364) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2366 = stablehlo.multiply %1749, %1749 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2367 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2368 = stablehlo.reduce(%2366 init: %2367) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2369 = stablehlo.multiply %1748, %1748 : tensor<1x2048xf32> loc(#loc5242)
    %2370 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2371 = stablehlo.reduce(%2369 init: %2370) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2372 = stablehlo.multiply %1747, %1747 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2373 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2374 = stablehlo.reduce(%2372 init: %2373) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2375 = stablehlo.multiply %1746, %1746 : tensor<1x2048xf32> loc(#loc5242)
    %2376 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2377 = stablehlo.reduce(%2375 init: %2376) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2378 = stablehlo.multiply %1745, %1745 : tensor<1x2048xf32> loc(#loc5242)
    %2379 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2380 = stablehlo.reduce(%2378 init: %2379) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2381 = stablehlo.multiply %1744, %1744 : tensor<1x2048xf32> loc(#loc5242)
    %2382 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2383 = stablehlo.reduce(%2381 init: %2382) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2384 = stablehlo.multiply %1743, %1743 : tensor<1x2048xf32> loc(#loc5242)
    %2385 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2386 = stablehlo.reduce(%2384 init: %2385) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2387 = stablehlo.multiply %1742, %1742 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2388 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2389 = stablehlo.reduce(%2387 init: %2388) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2390 = stablehlo.multiply %1741, %1741 : tensor<1x64xf32> loc(#loc5238)
    %2391 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2392 = stablehlo.reduce(%2390 init: %2391) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2393 = stablehlo.multiply %1740, %1740 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2394 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2395 = stablehlo.reduce(%2393 init: %2394) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2396 = stablehlo.multiply %1739, %1739 : tensor<1x8192xf32> loc(#loc5238)
    %2397 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2398 = stablehlo.reduce(%2396 init: %2397) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2399 = stablehlo.multiply %1738, %1738 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2400 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2401 = stablehlo.reduce(%2399 init: %2400) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2402 = stablehlo.multiply %1737, %1737 : tensor<1x2048xf32> loc(#loc5242)
    %2403 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2404 = stablehlo.reduce(%2402 init: %2403) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2405 = stablehlo.multiply %1736, %1736 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2406 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2407 = stablehlo.reduce(%2405 init: %2406) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2408 = stablehlo.multiply %1735, %1735 : tensor<1x2048xf32> loc(#loc5242)
    %2409 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2410 = stablehlo.reduce(%2408 init: %2409) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2411 = stablehlo.multiply %1734, %1734 : tensor<1x2048xf32> loc(#loc5242)
    %2412 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2413 = stablehlo.reduce(%2411 init: %2412) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2414 = stablehlo.multiply %1733, %1733 : tensor<1x2048xf32> loc(#loc5242)
    %2415 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2416 = stablehlo.reduce(%2414 init: %2415) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2417 = stablehlo.multiply %1732, %1732 : tensor<1x2048xf32> loc(#loc5242)
    %2418 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2419 = stablehlo.reduce(%2417 init: %2418) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2420 = stablehlo.multiply %1731, %1731 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2421 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2422 = stablehlo.reduce(%2420 init: %2421) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2423 = stablehlo.multiply %1730, %1730 : tensor<1x64xf32> loc(#loc5238)
    %2424 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2425 = stablehlo.reduce(%2423 init: %2424) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2426 = stablehlo.multiply %1729, %1729 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2427 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2428 = stablehlo.reduce(%2426 init: %2427) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2429 = stablehlo.multiply %1728, %1728 : tensor<1x8192xf32> loc(#loc5238)
    %2430 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2431 = stablehlo.reduce(%2429 init: %2430) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2432 = stablehlo.multiply %1727, %1727 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2433 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2434 = stablehlo.reduce(%2432 init: %2433) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2435 = stablehlo.multiply %1726, %1726 : tensor<1x2048xf32> loc(#loc5242)
    %2436 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2437 = stablehlo.reduce(%2435 init: %2436) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2438 = stablehlo.multiply %1725, %1725 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2439 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2440 = stablehlo.reduce(%2438 init: %2439) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2441 = stablehlo.multiply %1724, %1724 : tensor<1x2048xf32> loc(#loc5242)
    %2442 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2443 = stablehlo.reduce(%2441 init: %2442) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2444 = stablehlo.multiply %1723, %1723 : tensor<1x2048xf32> loc(#loc5242)
    %2445 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2446 = stablehlo.reduce(%2444 init: %2445) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2447 = stablehlo.multiply %1722, %1722 : tensor<1x2048xf32> loc(#loc5242)
    %2448 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2449 = stablehlo.reduce(%2447 init: %2448) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2450 = stablehlo.multiply %1721, %1721 : tensor<1x2048xf32> loc(#loc5242)
    %2451 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2452 = stablehlo.reduce(%2450 init: %2451) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2453 = stablehlo.multiply %1720, %1720 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2454 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2455 = stablehlo.reduce(%2453 init: %2454) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2456 = stablehlo.multiply %1719, %1719 : tensor<1x64xf32> loc(#loc5238)
    %2457 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2458 = stablehlo.reduce(%2456 init: %2457) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2459 = stablehlo.multiply %1718, %1718 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2460 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2461 = stablehlo.reduce(%2459 init: %2460) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2462 = stablehlo.multiply %1717, %1717 : tensor<1x8192xf32> loc(#loc5238)
    %2463 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2464 = stablehlo.reduce(%2462 init: %2463) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2465 = stablehlo.multiply %1716, %1716 : tensor<1x2048x8192xf32> loc(#loc5238)
    %2466 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2467 = stablehlo.reduce(%2465 init: %2466) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2468 = stablehlo.multiply %1715, %1715 : tensor<1x2048xf32> loc(#loc5242)
    %2469 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2470 = stablehlo.reduce(%2468 init: %2469) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2471 = stablehlo.multiply %1714, %1714 : tensor<1x8192x2048xf32> loc(#loc5238)
    %2472 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %2473 = stablehlo.reduce(%2471 init: %2472) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %2474 = stablehlo.multiply %1713, %1713 : tensor<1x2048xf32> loc(#loc5242)
    %2475 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2476 = stablehlo.reduce(%2474 init: %2475) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2477 = stablehlo.multiply %1712, %1712 : tensor<1x2048xf32> loc(#loc5242)
    %2478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2479 = stablehlo.reduce(%2477 init: %2478) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2480 = stablehlo.multiply %1711, %1711 : tensor<1x2048xf32> loc(#loc5242)
    %2481 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2482 = stablehlo.reduce(%2480 init: %2481) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2483 = stablehlo.multiply %1710, %1710 : tensor<1x2048xf32> loc(#loc5242)
    %2484 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %2485 = stablehlo.reduce(%2483 init: %2484) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %2486 = stablehlo.multiply %1709, %1709 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %2487 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %2488 = stablehlo.reduce(%2486 init: %2487) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %2489 = stablehlo.multiply %1708, %1708 : tensor<1x64xf32> loc(#loc5238)
    %2490 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %2491 = stablehlo.reduce(%2489 init: %2490) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %2492 = stablehlo.multiply %1707, %1707 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %2493 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %2494 = stablehlo.reduce(%2492 init: %2493) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %2495 = stablehlo.broadcast_in_dim %1954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2496 = stablehlo.broadcast_in_dim %1957, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2497 = stablehlo.broadcast_in_dim %1960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2498 = stablehlo.broadcast_in_dim %1963, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2499 = stablehlo.broadcast_in_dim %1966, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2500 = stablehlo.broadcast_in_dim %1969, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2501 = stablehlo.broadcast_in_dim %1972, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2502 = stablehlo.broadcast_in_dim %1975, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2503 = stablehlo.broadcast_in_dim %1978, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2504 = stablehlo.broadcast_in_dim %1981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2505 = stablehlo.broadcast_in_dim %1984, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2506 = stablehlo.broadcast_in_dim %1987, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2507 = stablehlo.broadcast_in_dim %1990, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2508 = stablehlo.broadcast_in_dim %1993, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2509 = stablehlo.broadcast_in_dim %1996, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2510 = stablehlo.broadcast_in_dim %1999, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2511 = stablehlo.broadcast_in_dim %2002, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2512 = stablehlo.broadcast_in_dim %2005, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2513 = stablehlo.broadcast_in_dim %2008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2514 = stablehlo.broadcast_in_dim %2011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2515 = stablehlo.broadcast_in_dim %2014, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2516 = stablehlo.broadcast_in_dim %2017, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2517 = stablehlo.broadcast_in_dim %2020, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2518 = stablehlo.broadcast_in_dim %2023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2519 = stablehlo.broadcast_in_dim %2026, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2520 = stablehlo.broadcast_in_dim %2029, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2521 = stablehlo.broadcast_in_dim %2032, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2522 = stablehlo.broadcast_in_dim %2035, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2523 = stablehlo.broadcast_in_dim %2038, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2524 = stablehlo.broadcast_in_dim %2041, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2525 = stablehlo.broadcast_in_dim %2044, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2526 = stablehlo.broadcast_in_dim %2047, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2527 = stablehlo.broadcast_in_dim %2050, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2528 = stablehlo.broadcast_in_dim %2053, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2529 = stablehlo.broadcast_in_dim %2056, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2530 = stablehlo.broadcast_in_dim %2059, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2531 = stablehlo.broadcast_in_dim %2062, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2532 = stablehlo.broadcast_in_dim %2065, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2533 = stablehlo.broadcast_in_dim %2068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2534 = stablehlo.broadcast_in_dim %2071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2535 = stablehlo.broadcast_in_dim %2074, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2536 = stablehlo.broadcast_in_dim %2077, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2537 = stablehlo.broadcast_in_dim %2080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2538 = stablehlo.broadcast_in_dim %2083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2539 = stablehlo.broadcast_in_dim %2086, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2540 = stablehlo.broadcast_in_dim %2089, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2541 = stablehlo.broadcast_in_dim %2092, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2542 = stablehlo.broadcast_in_dim %2095, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2543 = stablehlo.broadcast_in_dim %2098, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2544 = stablehlo.broadcast_in_dim %2101, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2545 = stablehlo.broadcast_in_dim %2104, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2546 = stablehlo.broadcast_in_dim %2107, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2547 = stablehlo.broadcast_in_dim %2110, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2548 = stablehlo.broadcast_in_dim %2113, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2549 = stablehlo.broadcast_in_dim %2116, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2550 = stablehlo.broadcast_in_dim %2119, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2551 = stablehlo.broadcast_in_dim %2122, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2552 = stablehlo.broadcast_in_dim %2125, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2553 = stablehlo.broadcast_in_dim %2128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2554 = stablehlo.broadcast_in_dim %2131, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2555 = stablehlo.broadcast_in_dim %2134, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2556 = stablehlo.broadcast_in_dim %2137, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2557 = stablehlo.broadcast_in_dim %2140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2558 = stablehlo.broadcast_in_dim %2143, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2559 = stablehlo.broadcast_in_dim %2146, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2560 = stablehlo.broadcast_in_dim %2149, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2561 = stablehlo.broadcast_in_dim %2152, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2562 = stablehlo.broadcast_in_dim %2155, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2563 = stablehlo.broadcast_in_dim %2158, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2564 = stablehlo.broadcast_in_dim %2161, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2565 = stablehlo.broadcast_in_dim %2164, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2566 = stablehlo.broadcast_in_dim %2167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2567 = stablehlo.broadcast_in_dim %2170, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2568 = stablehlo.broadcast_in_dim %2173, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2569 = stablehlo.broadcast_in_dim %2176, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2570 = stablehlo.broadcast_in_dim %2179, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2571 = stablehlo.broadcast_in_dim %2182, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2572 = stablehlo.broadcast_in_dim %2185, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2573 = stablehlo.broadcast_in_dim %2188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2574 = stablehlo.broadcast_in_dim %2191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2575 = stablehlo.broadcast_in_dim %2194, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2576 = stablehlo.broadcast_in_dim %2197, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2577 = stablehlo.broadcast_in_dim %2200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2578 = stablehlo.broadcast_in_dim %2203, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2579 = stablehlo.broadcast_in_dim %2206, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2580 = stablehlo.broadcast_in_dim %2209, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2581 = stablehlo.broadcast_in_dim %2212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2582 = stablehlo.broadcast_in_dim %2215, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2583 = stablehlo.broadcast_in_dim %2218, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2584 = stablehlo.broadcast_in_dim %2221, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2585 = stablehlo.broadcast_in_dim %2224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2586 = stablehlo.broadcast_in_dim %2227, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2587 = stablehlo.broadcast_in_dim %2230, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2588 = stablehlo.broadcast_in_dim %2233, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2589 = stablehlo.broadcast_in_dim %2236, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2590 = stablehlo.broadcast_in_dim %2239, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2591 = stablehlo.broadcast_in_dim %2242, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2592 = stablehlo.broadcast_in_dim %2245, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2593 = stablehlo.broadcast_in_dim %2248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2594 = stablehlo.broadcast_in_dim %2251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2595 = stablehlo.broadcast_in_dim %2254, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2596 = stablehlo.broadcast_in_dim %2257, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2597 = stablehlo.broadcast_in_dim %2260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2598 = stablehlo.broadcast_in_dim %2263, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2599 = stablehlo.broadcast_in_dim %2266, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2600 = stablehlo.broadcast_in_dim %2269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2601 = stablehlo.broadcast_in_dim %2272, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2602 = stablehlo.broadcast_in_dim %2275, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2603 = stablehlo.broadcast_in_dim %2278, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2604 = stablehlo.broadcast_in_dim %2281, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2605 = stablehlo.broadcast_in_dim %2284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2606 = stablehlo.broadcast_in_dim %2287, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2607 = stablehlo.broadcast_in_dim %2290, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2608 = stablehlo.broadcast_in_dim %2293, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2609 = stablehlo.broadcast_in_dim %2296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2610 = stablehlo.broadcast_in_dim %2299, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2611 = stablehlo.broadcast_in_dim %2302, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2612 = stablehlo.broadcast_in_dim %2305, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2613 = stablehlo.broadcast_in_dim %2308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2614 = stablehlo.broadcast_in_dim %2311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2615 = stablehlo.broadcast_in_dim %2314, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2616 = stablehlo.broadcast_in_dim %2317, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2617 = stablehlo.broadcast_in_dim %2320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2618 = stablehlo.broadcast_in_dim %2323, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2619 = stablehlo.broadcast_in_dim %2326, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2620 = stablehlo.broadcast_in_dim %2329, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2621 = stablehlo.broadcast_in_dim %2332, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2622 = stablehlo.broadcast_in_dim %2335, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2623 = stablehlo.broadcast_in_dim %2338, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2624 = stablehlo.broadcast_in_dim %2341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2625 = stablehlo.broadcast_in_dim %2344, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2626 = stablehlo.broadcast_in_dim %2347, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2627 = stablehlo.broadcast_in_dim %2350, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2628 = stablehlo.broadcast_in_dim %2353, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2629 = stablehlo.broadcast_in_dim %2356, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2630 = stablehlo.broadcast_in_dim %2359, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2631 = stablehlo.broadcast_in_dim %2362, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2632 = stablehlo.broadcast_in_dim %2365, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2633 = stablehlo.broadcast_in_dim %2368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2634 = stablehlo.broadcast_in_dim %2371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2635 = stablehlo.broadcast_in_dim %2374, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2636 = stablehlo.broadcast_in_dim %2377, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2637 = stablehlo.broadcast_in_dim %2380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2638 = stablehlo.broadcast_in_dim %2383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2639 = stablehlo.broadcast_in_dim %2386, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2640 = stablehlo.broadcast_in_dim %2389, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2641 = stablehlo.broadcast_in_dim %2392, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2642 = stablehlo.broadcast_in_dim %2395, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2643 = stablehlo.broadcast_in_dim %2398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2644 = stablehlo.broadcast_in_dim %2401, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2645 = stablehlo.broadcast_in_dim %2404, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2646 = stablehlo.broadcast_in_dim %2407, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2647 = stablehlo.broadcast_in_dim %2410, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2648 = stablehlo.broadcast_in_dim %2413, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2649 = stablehlo.broadcast_in_dim %2416, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2650 = stablehlo.broadcast_in_dim %2419, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2651 = stablehlo.broadcast_in_dim %2422, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2652 = stablehlo.broadcast_in_dim %2425, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2653 = stablehlo.broadcast_in_dim %2428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2654 = stablehlo.broadcast_in_dim %2431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2655 = stablehlo.broadcast_in_dim %2434, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2656 = stablehlo.broadcast_in_dim %2437, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2657 = stablehlo.broadcast_in_dim %2440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2658 = stablehlo.broadcast_in_dim %2443, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2659 = stablehlo.broadcast_in_dim %2446, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2660 = stablehlo.broadcast_in_dim %2449, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2661 = stablehlo.broadcast_in_dim %2452, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2662 = stablehlo.broadcast_in_dim %2455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2663 = stablehlo.broadcast_in_dim %2458, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2664 = stablehlo.broadcast_in_dim %2461, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2665 = stablehlo.broadcast_in_dim %2464, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2666 = stablehlo.broadcast_in_dim %2467, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2667 = stablehlo.broadcast_in_dim %2470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2668 = stablehlo.broadcast_in_dim %2473, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2669 = stablehlo.broadcast_in_dim %2476, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2670 = stablehlo.broadcast_in_dim %2479, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2671 = stablehlo.broadcast_in_dim %2482, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2672 = stablehlo.broadcast_in_dim %2485, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2673 = stablehlo.broadcast_in_dim %2488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2674 = stablehlo.broadcast_in_dim %2491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2675 = stablehlo.broadcast_in_dim %2494, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %2676 = stablehlo.concatenate %2495, %2496, %2497, %2498, %2499, %2500, %2501, %2502, %2503, %2504, %2505, %2506, %2507, %2508, %2509, %2510, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2677 = stablehlo.concatenate %2511, %2512, %2513, %2514, %2515, %2516, %2517, %2518, %2519, %2520, %2521, %2522, %2523, %2524, %2525, %2526, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2678 = stablehlo.concatenate %2527, %2528, %2529, %2530, %2531, %2532, %2533, %2534, %2535, %2536, %2537, %2538, %2539, %2540, %2541, %2542, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2679 = stablehlo.concatenate %2543, %2544, %2545, %2546, %2547, %2548, %2549, %2550, %2551, %2552, %2553, %2554, %2555, %2556, %2557, %2558, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2680 = stablehlo.concatenate %2559, %2560, %2561, %2562, %2563, %2564, %2565, %2566, %2567, %2568, %2569, %2570, %2571, %2572, %2573, %2574, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2681 = stablehlo.concatenate %2575, %2576, %2577, %2578, %2579, %2580, %2581, %2582, %2583, %2584, %2585, %2586, %2587, %2588, %2589, %2590, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2682 = stablehlo.concatenate %2591, %2592, %2593, %2594, %2595, %2596, %2597, %2598, %2599, %2600, %2601, %2602, %2603, %2604, %2605, %2606, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2683 = stablehlo.concatenate %2607, %2608, %2609, %2610, %2611, %2612, %2613, %2614, %2615, %2616, %2617, %2618, %2619, %2620, %2621, %2622, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2684 = stablehlo.concatenate %2623, %2624, %2625, %2626, %2627, %2628, %2629, %2630, %2631, %2632, %2633, %2634, %2635, %2636, %2637, %2638, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2685 = stablehlo.concatenate %2639, %2640, %2641, %2642, %2643, %2644, %2645, %2646, %2647, %2648, %2649, %2650, %2651, %2652, %2653, %2654, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2686 = stablehlo.concatenate %2655, %2656, %2657, %2658, %2659, %2660, %2661, %2662, %2663, %2664, %2665, %2666, %2667, %2668, %2669, %2670, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %2687 = stablehlo.concatenate %2671, %2672, %2673, %2674, %2675, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5xf32> loc(#loc5247)
    %2688 = stablehlo.concatenate %2676, %2677, %2678, %2679, %2680, %2681, %2682, %2683, %2684, %2685, %2686, %2687, dim = 0 : (tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<5xf32>) -> tensor<181xf32> loc(#loc5247)
    %2689 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5248)
    %2690 = stablehlo.reduce(%2688 init: %2689) across dimensions = [0] : (tensor<181xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5248)
      stablehlo.return %29929 : tensor<f32> loc(#loc5248)
    } loc(#loc5248)
    %2691 = stablehlo.sqrt %2690 : tensor<f32> loc(#loc5249)
    %2692 = call @isfinite(%2691) : (tensor<f32>) -> tensor<i1> loc(#loc5250)
    %2693 = stablehlo.constant dense<true> : tensor<i1> loc(#loc5251)
    %2694 = stablehlo.reduce(%2692 init: %2693) across dimensions = [] : (tensor<i1>, tensor<i1>) -> tensor<i1>
     reducer(%arg558: tensor<i1> loc(unknown), %arg559: tensor<i1> loc(unknown))  {
      %29929 = stablehlo.and %arg558, %arg559 : tensor<i1> loc(#loc5251)
      stablehlo.return %29929 : tensor<i1> loc(#loc5251)
    } loc(#loc5251)
    %2695 = stablehlo.convert %2694 : (tensor<i1>) -> tensor<f32> loc(#loc5252)
    %2696 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %2697 = stablehlo.divide %2696, %2691 : tensor<f32> loc(#loc5253)
    %2698 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %2699 = stablehlo.minimum %2698, %2697 : tensor<f32> loc(#loc5254)
    %2700 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5255)
    %2701 = stablehlo.multiply %1198, %2700 : tensor<2048xf32> loc(#loc5255)
    %2702 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5255)
    %2703 = stablehlo.multiply %1207, %2702 : tensor<2048xf32> loc(#loc5255)
    %2704 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5255)
    %2705 = stablehlo.multiply %1895, %2704 : tensor<2048x2048xf32> loc(#loc5255)
    %2706 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5255)
    %2707 = stablehlo.multiply %1190, %2706 : tensor<51200xf32> loc(#loc5255)
    %2708 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5255)
    %2709 = stablehlo.multiply %1903, %2708 : tensor<2048x51200xf32> loc(#loc5255)
    %2710 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2711 = stablehlo.multiply %1882, %2710 : tensor<1x8192xf32> loc(#loc5255)
    %2712 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2713 = stablehlo.multiply %1881, %2712 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2714 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2715 = stablehlo.multiply %1880, %2714 : tensor<1x2048xf32> loc(#loc5255)
    %2716 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2717 = stablehlo.multiply %1879, %2716 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2718 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2719 = stablehlo.multiply %1878, %2718 : tensor<1x2048xf32> loc(#loc5255)
    %2720 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2721 = stablehlo.multiply %1877, %2720 : tensor<1x2048xf32> loc(#loc5255)
    %2722 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2723 = stablehlo.multiply %1876, %2722 : tensor<1x2048xf32> loc(#loc5255)
    %2724 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2725 = stablehlo.multiply %1875, %2724 : tensor<1x2048xf32> loc(#loc5255)
    %2726 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2727 = stablehlo.multiply %1874, %2726 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2728 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2729 = stablehlo.multiply %1873, %2728 : tensor<1x64xf32> loc(#loc5255)
    %2730 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2731 = stablehlo.multiply %1872, %2730 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2732 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2733 = stablehlo.multiply %1871, %2732 : tensor<1x8192xf32> loc(#loc5255)
    %2734 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2735 = stablehlo.multiply %1870, %2734 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2736 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2737 = stablehlo.multiply %1869, %2736 : tensor<1x2048xf32> loc(#loc5255)
    %2738 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2739 = stablehlo.multiply %1868, %2738 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2740 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2741 = stablehlo.multiply %1867, %2740 : tensor<1x2048xf32> loc(#loc5255)
    %2742 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2743 = stablehlo.multiply %1866, %2742 : tensor<1x2048xf32> loc(#loc5255)
    %2744 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2745 = stablehlo.multiply %1865, %2744 : tensor<1x2048xf32> loc(#loc5255)
    %2746 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2747 = stablehlo.multiply %1864, %2746 : tensor<1x2048xf32> loc(#loc5255)
    %2748 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2749 = stablehlo.multiply %1863, %2748 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2750 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2751 = stablehlo.multiply %1862, %2750 : tensor<1x64xf32> loc(#loc5255)
    %2752 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2753 = stablehlo.multiply %1861, %2752 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2754 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2755 = stablehlo.multiply %1860, %2754 : tensor<1x8192xf32> loc(#loc5255)
    %2756 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2757 = stablehlo.multiply %1859, %2756 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2758 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2759 = stablehlo.multiply %1858, %2758 : tensor<1x2048xf32> loc(#loc5255)
    %2760 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2761 = stablehlo.multiply %1857, %2760 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2762 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2763 = stablehlo.multiply %1856, %2762 : tensor<1x2048xf32> loc(#loc5255)
    %2764 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2765 = stablehlo.multiply %1855, %2764 : tensor<1x2048xf32> loc(#loc5255)
    %2766 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2767 = stablehlo.multiply %1854, %2766 : tensor<1x2048xf32> loc(#loc5255)
    %2768 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2769 = stablehlo.multiply %1853, %2768 : tensor<1x2048xf32> loc(#loc5255)
    %2770 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2771 = stablehlo.multiply %1852, %2770 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2772 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2773 = stablehlo.multiply %1851, %2772 : tensor<1x64xf32> loc(#loc5255)
    %2774 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2775 = stablehlo.multiply %1850, %2774 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2776 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2777 = stablehlo.multiply %1849, %2776 : tensor<1x8192xf32> loc(#loc5255)
    %2778 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2779 = stablehlo.multiply %1848, %2778 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2780 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2781 = stablehlo.multiply %1847, %2780 : tensor<1x2048xf32> loc(#loc5255)
    %2782 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2783 = stablehlo.multiply %1846, %2782 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2784 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2785 = stablehlo.multiply %1845, %2784 : tensor<1x2048xf32> loc(#loc5255)
    %2786 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2787 = stablehlo.multiply %1844, %2786 : tensor<1x2048xf32> loc(#loc5255)
    %2788 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2789 = stablehlo.multiply %1843, %2788 : tensor<1x2048xf32> loc(#loc5255)
    %2790 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2791 = stablehlo.multiply %1842, %2790 : tensor<1x2048xf32> loc(#loc5255)
    %2792 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2793 = stablehlo.multiply %1841, %2792 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2794 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2795 = stablehlo.multiply %1840, %2794 : tensor<1x64xf32> loc(#loc5255)
    %2796 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2797 = stablehlo.multiply %1839, %2796 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2798 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2799 = stablehlo.multiply %1838, %2798 : tensor<1x8192xf32> loc(#loc5255)
    %2800 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2801 = stablehlo.multiply %1837, %2800 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2802 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2803 = stablehlo.multiply %1836, %2802 : tensor<1x2048xf32> loc(#loc5255)
    %2804 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2805 = stablehlo.multiply %1835, %2804 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2806 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2807 = stablehlo.multiply %1834, %2806 : tensor<1x2048xf32> loc(#loc5255)
    %2808 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2809 = stablehlo.multiply %1833, %2808 : tensor<1x2048xf32> loc(#loc5255)
    %2810 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2811 = stablehlo.multiply %1832, %2810 : tensor<1x2048xf32> loc(#loc5255)
    %2812 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2813 = stablehlo.multiply %1831, %2812 : tensor<1x2048xf32> loc(#loc5255)
    %2814 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2815 = stablehlo.multiply %1830, %2814 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2816 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2817 = stablehlo.multiply %1829, %2816 : tensor<1x64xf32> loc(#loc5255)
    %2818 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2819 = stablehlo.multiply %1828, %2818 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2820 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2821 = stablehlo.multiply %1827, %2820 : tensor<1x8192xf32> loc(#loc5255)
    %2822 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2823 = stablehlo.multiply %1826, %2822 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2824 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2825 = stablehlo.multiply %1825, %2824 : tensor<1x2048xf32> loc(#loc5255)
    %2826 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2827 = stablehlo.multiply %1824, %2826 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2828 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2829 = stablehlo.multiply %1823, %2828 : tensor<1x2048xf32> loc(#loc5255)
    %2830 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2831 = stablehlo.multiply %1822, %2830 : tensor<1x2048xf32> loc(#loc5255)
    %2832 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2833 = stablehlo.multiply %1821, %2832 : tensor<1x2048xf32> loc(#loc5255)
    %2834 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2835 = stablehlo.multiply %1820, %2834 : tensor<1x2048xf32> loc(#loc5255)
    %2836 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2837 = stablehlo.multiply %1819, %2836 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2838 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2839 = stablehlo.multiply %1818, %2838 : tensor<1x64xf32> loc(#loc5255)
    %2840 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2841 = stablehlo.multiply %1817, %2840 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2842 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2843 = stablehlo.multiply %1816, %2842 : tensor<1x8192xf32> loc(#loc5255)
    %2844 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2845 = stablehlo.multiply %1815, %2844 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2846 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2847 = stablehlo.multiply %1814, %2846 : tensor<1x2048xf32> loc(#loc5255)
    %2848 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2849 = stablehlo.multiply %1813, %2848 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2850 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2851 = stablehlo.multiply %1812, %2850 : tensor<1x2048xf32> loc(#loc5255)
    %2852 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2853 = stablehlo.multiply %1811, %2852 : tensor<1x2048xf32> loc(#loc5255)
    %2854 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2855 = stablehlo.multiply %1810, %2854 : tensor<1x2048xf32> loc(#loc5255)
    %2856 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2857 = stablehlo.multiply %1809, %2856 : tensor<1x2048xf32> loc(#loc5255)
    %2858 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2859 = stablehlo.multiply %1808, %2858 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2860 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2861 = stablehlo.multiply %1807, %2860 : tensor<1x64xf32> loc(#loc5255)
    %2862 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2863 = stablehlo.multiply %1806, %2862 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2864 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2865 = stablehlo.multiply %1805, %2864 : tensor<1x8192xf32> loc(#loc5255)
    %2866 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2867 = stablehlo.multiply %1804, %2866 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2868 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2869 = stablehlo.multiply %1803, %2868 : tensor<1x2048xf32> loc(#loc5255)
    %2870 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2871 = stablehlo.multiply %1802, %2870 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2872 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2873 = stablehlo.multiply %1801, %2872 : tensor<1x2048xf32> loc(#loc5255)
    %2874 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2875 = stablehlo.multiply %1800, %2874 : tensor<1x2048xf32> loc(#loc5255)
    %2876 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2877 = stablehlo.multiply %1799, %2876 : tensor<1x2048xf32> loc(#loc5255)
    %2878 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2879 = stablehlo.multiply %1798, %2878 : tensor<1x2048xf32> loc(#loc5255)
    %2880 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2881 = stablehlo.multiply %1797, %2880 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2882 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2883 = stablehlo.multiply %1796, %2882 : tensor<1x64xf32> loc(#loc5255)
    %2884 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2885 = stablehlo.multiply %1795, %2884 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2886 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2887 = stablehlo.multiply %1794, %2886 : tensor<1x8192xf32> loc(#loc5255)
    %2888 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2889 = stablehlo.multiply %1793, %2888 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2890 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2891 = stablehlo.multiply %1792, %2890 : tensor<1x2048xf32> loc(#loc5255)
    %2892 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2893 = stablehlo.multiply %1791, %2892 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2894 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2895 = stablehlo.multiply %1790, %2894 : tensor<1x2048xf32> loc(#loc5255)
    %2896 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2897 = stablehlo.multiply %1789, %2896 : tensor<1x2048xf32> loc(#loc5255)
    %2898 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2899 = stablehlo.multiply %1788, %2898 : tensor<1x2048xf32> loc(#loc5255)
    %2900 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2901 = stablehlo.multiply %1787, %2900 : tensor<1x2048xf32> loc(#loc5255)
    %2902 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2903 = stablehlo.multiply %1786, %2902 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2904 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2905 = stablehlo.multiply %1785, %2904 : tensor<1x64xf32> loc(#loc5255)
    %2906 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2907 = stablehlo.multiply %1784, %2906 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2908 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2909 = stablehlo.multiply %1783, %2908 : tensor<1x8192xf32> loc(#loc5255)
    %2910 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2911 = stablehlo.multiply %1782, %2910 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2912 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2913 = stablehlo.multiply %1781, %2912 : tensor<1x2048xf32> loc(#loc5255)
    %2914 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2915 = stablehlo.multiply %1780, %2914 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2916 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2917 = stablehlo.multiply %1779, %2916 : tensor<1x2048xf32> loc(#loc5255)
    %2918 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2919 = stablehlo.multiply %1778, %2918 : tensor<1x2048xf32> loc(#loc5255)
    %2920 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2921 = stablehlo.multiply %1777, %2920 : tensor<1x2048xf32> loc(#loc5255)
    %2922 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2923 = stablehlo.multiply %1776, %2922 : tensor<1x2048xf32> loc(#loc5255)
    %2924 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2925 = stablehlo.multiply %1775, %2924 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2926 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2927 = stablehlo.multiply %1774, %2926 : tensor<1x64xf32> loc(#loc5255)
    %2928 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2929 = stablehlo.multiply %1773, %2928 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2930 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2931 = stablehlo.multiply %1772, %2930 : tensor<1x8192xf32> loc(#loc5255)
    %2932 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2933 = stablehlo.multiply %1771, %2932 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2934 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2935 = stablehlo.multiply %1770, %2934 : tensor<1x2048xf32> loc(#loc5255)
    %2936 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2937 = stablehlo.multiply %1769, %2936 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2938 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2939 = stablehlo.multiply %1768, %2938 : tensor<1x2048xf32> loc(#loc5255)
    %2940 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2941 = stablehlo.multiply %1767, %2940 : tensor<1x2048xf32> loc(#loc5255)
    %2942 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2943 = stablehlo.multiply %1766, %2942 : tensor<1x2048xf32> loc(#loc5255)
    %2944 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2945 = stablehlo.multiply %1765, %2944 : tensor<1x2048xf32> loc(#loc5255)
    %2946 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2947 = stablehlo.multiply %1764, %2946 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2948 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2949 = stablehlo.multiply %1763, %2948 : tensor<1x64xf32> loc(#loc5255)
    %2950 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2951 = stablehlo.multiply %1762, %2950 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2952 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2953 = stablehlo.multiply %1761, %2952 : tensor<1x8192xf32> loc(#loc5255)
    %2954 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2955 = stablehlo.multiply %1760, %2954 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2956 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2957 = stablehlo.multiply %1759, %2956 : tensor<1x2048xf32> loc(#loc5255)
    %2958 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2959 = stablehlo.multiply %1758, %2958 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2960 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2961 = stablehlo.multiply %1757, %2960 : tensor<1x2048xf32> loc(#loc5255)
    %2962 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2963 = stablehlo.multiply %1756, %2962 : tensor<1x2048xf32> loc(#loc5255)
    %2964 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2965 = stablehlo.multiply %1755, %2964 : tensor<1x2048xf32> loc(#loc5255)
    %2966 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2967 = stablehlo.multiply %1754, %2966 : tensor<1x2048xf32> loc(#loc5255)
    %2968 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2969 = stablehlo.multiply %1753, %2968 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2970 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2971 = stablehlo.multiply %1752, %2970 : tensor<1x64xf32> loc(#loc5255)
    %2972 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2973 = stablehlo.multiply %1751, %2972 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2974 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2975 = stablehlo.multiply %1750, %2974 : tensor<1x8192xf32> loc(#loc5255)
    %2976 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2977 = stablehlo.multiply %1749, %2976 : tensor<1x2048x8192xf32> loc(#loc5255)
    %2978 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2979 = stablehlo.multiply %1748, %2978 : tensor<1x2048xf32> loc(#loc5255)
    %2980 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %2981 = stablehlo.multiply %1747, %2980 : tensor<1x8192x2048xf32> loc(#loc5255)
    %2982 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2983 = stablehlo.multiply %1746, %2982 : tensor<1x2048xf32> loc(#loc5255)
    %2984 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2985 = stablehlo.multiply %1745, %2984 : tensor<1x2048xf32> loc(#loc5255)
    %2986 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2987 = stablehlo.multiply %1744, %2986 : tensor<1x2048xf32> loc(#loc5255)
    %2988 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %2989 = stablehlo.multiply %1743, %2988 : tensor<1x2048xf32> loc(#loc5255)
    %2990 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2991 = stablehlo.multiply %1742, %2990 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %2992 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %2993 = stablehlo.multiply %1741, %2992 : tensor<1x64xf32> loc(#loc5255)
    %2994 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %2995 = stablehlo.multiply %1740, %2994 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %2996 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %2997 = stablehlo.multiply %1739, %2996 : tensor<1x8192xf32> loc(#loc5255)
    %2998 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %2999 = stablehlo.multiply %1738, %2998 : tensor<1x2048x8192xf32> loc(#loc5255)
    %3000 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3001 = stablehlo.multiply %1737, %3000 : tensor<1x2048xf32> loc(#loc5255)
    %3002 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %3003 = stablehlo.multiply %1736, %3002 : tensor<1x8192x2048xf32> loc(#loc5255)
    %3004 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3005 = stablehlo.multiply %1735, %3004 : tensor<1x2048xf32> loc(#loc5255)
    %3006 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3007 = stablehlo.multiply %1734, %3006 : tensor<1x2048xf32> loc(#loc5255)
    %3008 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3009 = stablehlo.multiply %1733, %3008 : tensor<1x2048xf32> loc(#loc5255)
    %3010 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3011 = stablehlo.multiply %1732, %3010 : tensor<1x2048xf32> loc(#loc5255)
    %3012 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %3013 = stablehlo.multiply %1731, %3012 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %3014 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %3015 = stablehlo.multiply %1730, %3014 : tensor<1x64xf32> loc(#loc5255)
    %3016 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %3017 = stablehlo.multiply %1729, %3016 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %3018 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %3019 = stablehlo.multiply %1728, %3018 : tensor<1x8192xf32> loc(#loc5255)
    %3020 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %3021 = stablehlo.multiply %1727, %3020 : tensor<1x2048x8192xf32> loc(#loc5255)
    %3022 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3023 = stablehlo.multiply %1726, %3022 : tensor<1x2048xf32> loc(#loc5255)
    %3024 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %3025 = stablehlo.multiply %1725, %3024 : tensor<1x8192x2048xf32> loc(#loc5255)
    %3026 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3027 = stablehlo.multiply %1724, %3026 : tensor<1x2048xf32> loc(#loc5255)
    %3028 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3029 = stablehlo.multiply %1723, %3028 : tensor<1x2048xf32> loc(#loc5255)
    %3030 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3031 = stablehlo.multiply %1722, %3030 : tensor<1x2048xf32> loc(#loc5255)
    %3032 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3033 = stablehlo.multiply %1721, %3032 : tensor<1x2048xf32> loc(#loc5255)
    %3034 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %3035 = stablehlo.multiply %1720, %3034 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %3036 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %3037 = stablehlo.multiply %1719, %3036 : tensor<1x64xf32> loc(#loc5255)
    %3038 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %3039 = stablehlo.multiply %1718, %3038 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %3040 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5255)
    %3041 = stablehlo.multiply %1717, %3040 : tensor<1x8192xf32> loc(#loc5255)
    %3042 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5255)
    %3043 = stablehlo.multiply %1716, %3042 : tensor<1x2048x8192xf32> loc(#loc5255)
    %3044 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3045 = stablehlo.multiply %1715, %3044 : tensor<1x2048xf32> loc(#loc5255)
    %3046 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5255)
    %3047 = stablehlo.multiply %1714, %3046 : tensor<1x8192x2048xf32> loc(#loc5255)
    %3048 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3049 = stablehlo.multiply %1713, %3048 : tensor<1x2048xf32> loc(#loc5255)
    %3050 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3051 = stablehlo.multiply %1712, %3050 : tensor<1x2048xf32> loc(#loc5255)
    %3052 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3053 = stablehlo.multiply %1711, %3052 : tensor<1x2048xf32> loc(#loc5255)
    %3054 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5255)
    %3055 = stablehlo.multiply %1710, %3054 : tensor<1x2048xf32> loc(#loc5255)
    %3056 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %3057 = stablehlo.multiply %1709, %3056 : tensor<1x3x2048x32x64xf32> loc(#loc5255)
    %3058 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5255)
    %3059 = stablehlo.multiply %1708, %3058 : tensor<1x64xf32> loc(#loc5255)
    %3060 = stablehlo.broadcast_in_dim %2699, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5255)
    %3061 = stablehlo.multiply %1707, %3060 : tensor<1x2048x32x64xf32> loc(#loc5255)
    %3062 = stablehlo.multiply %2701, %2701 : tensor<2048xf32> loc(#loc5238)
    %3063 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %3064 = stablehlo.reduce(%3062 init: %3063) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %3065 = stablehlo.multiply %2703, %2703 : tensor<2048xf32> loc(#loc5238)
    %3066 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %3067 = stablehlo.reduce(%3065 init: %3066) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %3068 = stablehlo.multiply %2705, %2705 : tensor<2048x2048xf32> loc(#loc5238)
    %3069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3070 = stablehlo.reduce(%3068 init: %3069) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3071 = stablehlo.multiply %2707, %2707 : tensor<51200xf32> loc(#loc5238)
    %3072 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %3073 = stablehlo.reduce(%3071 init: %3072) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %3074 = stablehlo.multiply %2709, %2709 : tensor<2048x51200xf32> loc(#loc5238)
    %3075 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3076 = stablehlo.reduce(%3074 init: %3075) across dimensions = [0, 1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3077 = stablehlo.multiply %2711, %2711 : tensor<1x8192xf32> loc(#loc5238)
    %3078 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3079 = stablehlo.reduce(%3077 init: %3078) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3080 = stablehlo.multiply %2713, %2713 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3081 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3082 = stablehlo.reduce(%3080 init: %3081) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3083 = stablehlo.multiply %2715, %2715 : tensor<1x2048xf32> loc(#loc5242)
    %3084 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3085 = stablehlo.reduce(%3083 init: %3084) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3086 = stablehlo.multiply %2717, %2717 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3087 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3088 = stablehlo.reduce(%3086 init: %3087) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3089 = stablehlo.multiply %2719, %2719 : tensor<1x2048xf32> loc(#loc5242)
    %3090 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3091 = stablehlo.reduce(%3089 init: %3090) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3092 = stablehlo.multiply %2721, %2721 : tensor<1x2048xf32> loc(#loc5242)
    %3093 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3094 = stablehlo.reduce(%3092 init: %3093) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3095 = stablehlo.multiply %2723, %2723 : tensor<1x2048xf32> loc(#loc5242)
    %3096 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3097 = stablehlo.reduce(%3095 init: %3096) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3098 = stablehlo.multiply %2725, %2725 : tensor<1x2048xf32> loc(#loc5242)
    %3099 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3100 = stablehlo.reduce(%3098 init: %3099) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3101 = stablehlo.multiply %2727, %2727 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3102 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3103 = stablehlo.reduce(%3101 init: %3102) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3104 = stablehlo.multiply %2729, %2729 : tensor<1x64xf32> loc(#loc5238)
    %3105 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3106 = stablehlo.reduce(%3104 init: %3105) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3107 = stablehlo.multiply %2731, %2731 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3108 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3109 = stablehlo.reduce(%3107 init: %3108) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3110 = stablehlo.multiply %2733, %2733 : tensor<1x8192xf32> loc(#loc5238)
    %3111 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3112 = stablehlo.reduce(%3110 init: %3111) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3113 = stablehlo.multiply %2735, %2735 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3114 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3115 = stablehlo.reduce(%3113 init: %3114) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3116 = stablehlo.multiply %2737, %2737 : tensor<1x2048xf32> loc(#loc5242)
    %3117 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3118 = stablehlo.reduce(%3116 init: %3117) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3119 = stablehlo.multiply %2739, %2739 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3120 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3121 = stablehlo.reduce(%3119 init: %3120) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3122 = stablehlo.multiply %2741, %2741 : tensor<1x2048xf32> loc(#loc5242)
    %3123 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3124 = stablehlo.reduce(%3122 init: %3123) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3125 = stablehlo.multiply %2743, %2743 : tensor<1x2048xf32> loc(#loc5242)
    %3126 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3127 = stablehlo.reduce(%3125 init: %3126) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3128 = stablehlo.multiply %2745, %2745 : tensor<1x2048xf32> loc(#loc5242)
    %3129 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3130 = stablehlo.reduce(%3128 init: %3129) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3131 = stablehlo.multiply %2747, %2747 : tensor<1x2048xf32> loc(#loc5242)
    %3132 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3133 = stablehlo.reduce(%3131 init: %3132) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3134 = stablehlo.multiply %2749, %2749 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3135 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3136 = stablehlo.reduce(%3134 init: %3135) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3137 = stablehlo.multiply %2751, %2751 : tensor<1x64xf32> loc(#loc5238)
    %3138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3139 = stablehlo.reduce(%3137 init: %3138) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3140 = stablehlo.multiply %2753, %2753 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3141 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3142 = stablehlo.reduce(%3140 init: %3141) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3143 = stablehlo.multiply %2755, %2755 : tensor<1x8192xf32> loc(#loc5238)
    %3144 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3145 = stablehlo.reduce(%3143 init: %3144) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3146 = stablehlo.multiply %2757, %2757 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3147 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3148 = stablehlo.reduce(%3146 init: %3147) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3149 = stablehlo.multiply %2759, %2759 : tensor<1x2048xf32> loc(#loc5242)
    %3150 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3151 = stablehlo.reduce(%3149 init: %3150) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3152 = stablehlo.multiply %2761, %2761 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3153 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3154 = stablehlo.reduce(%3152 init: %3153) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3155 = stablehlo.multiply %2763, %2763 : tensor<1x2048xf32> loc(#loc5242)
    %3156 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3157 = stablehlo.reduce(%3155 init: %3156) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3158 = stablehlo.multiply %2765, %2765 : tensor<1x2048xf32> loc(#loc5242)
    %3159 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3160 = stablehlo.reduce(%3158 init: %3159) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3161 = stablehlo.multiply %2767, %2767 : tensor<1x2048xf32> loc(#loc5242)
    %3162 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3163 = stablehlo.reduce(%3161 init: %3162) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3164 = stablehlo.multiply %2769, %2769 : tensor<1x2048xf32> loc(#loc5242)
    %3165 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3166 = stablehlo.reduce(%3164 init: %3165) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3167 = stablehlo.multiply %2771, %2771 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3168 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3169 = stablehlo.reduce(%3167 init: %3168) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3170 = stablehlo.multiply %2773, %2773 : tensor<1x64xf32> loc(#loc5238)
    %3171 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3172 = stablehlo.reduce(%3170 init: %3171) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3173 = stablehlo.multiply %2775, %2775 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3174 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3175 = stablehlo.reduce(%3173 init: %3174) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3176 = stablehlo.multiply %2777, %2777 : tensor<1x8192xf32> loc(#loc5238)
    %3177 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3178 = stablehlo.reduce(%3176 init: %3177) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3179 = stablehlo.multiply %2779, %2779 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3180 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3181 = stablehlo.reduce(%3179 init: %3180) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3182 = stablehlo.multiply %2781, %2781 : tensor<1x2048xf32> loc(#loc5242)
    %3183 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3184 = stablehlo.reduce(%3182 init: %3183) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3185 = stablehlo.multiply %2783, %2783 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3186 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3187 = stablehlo.reduce(%3185 init: %3186) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3188 = stablehlo.multiply %2785, %2785 : tensor<1x2048xf32> loc(#loc5242)
    %3189 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3190 = stablehlo.reduce(%3188 init: %3189) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3191 = stablehlo.multiply %2787, %2787 : tensor<1x2048xf32> loc(#loc5242)
    %3192 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3193 = stablehlo.reduce(%3191 init: %3192) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3194 = stablehlo.multiply %2789, %2789 : tensor<1x2048xf32> loc(#loc5242)
    %3195 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3196 = stablehlo.reduce(%3194 init: %3195) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3197 = stablehlo.multiply %2791, %2791 : tensor<1x2048xf32> loc(#loc5242)
    %3198 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3199 = stablehlo.reduce(%3197 init: %3198) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3200 = stablehlo.multiply %2793, %2793 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3201 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3202 = stablehlo.reduce(%3200 init: %3201) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3203 = stablehlo.multiply %2795, %2795 : tensor<1x64xf32> loc(#loc5238)
    %3204 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3205 = stablehlo.reduce(%3203 init: %3204) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3206 = stablehlo.multiply %2797, %2797 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3207 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3208 = stablehlo.reduce(%3206 init: %3207) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3209 = stablehlo.multiply %2799, %2799 : tensor<1x8192xf32> loc(#loc5238)
    %3210 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3211 = stablehlo.reduce(%3209 init: %3210) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3212 = stablehlo.multiply %2801, %2801 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3213 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3214 = stablehlo.reduce(%3212 init: %3213) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3215 = stablehlo.multiply %2803, %2803 : tensor<1x2048xf32> loc(#loc5242)
    %3216 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3217 = stablehlo.reduce(%3215 init: %3216) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3218 = stablehlo.multiply %2805, %2805 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3219 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3220 = stablehlo.reduce(%3218 init: %3219) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3221 = stablehlo.multiply %2807, %2807 : tensor<1x2048xf32> loc(#loc5242)
    %3222 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3223 = stablehlo.reduce(%3221 init: %3222) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3224 = stablehlo.multiply %2809, %2809 : tensor<1x2048xf32> loc(#loc5242)
    %3225 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3226 = stablehlo.reduce(%3224 init: %3225) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3227 = stablehlo.multiply %2811, %2811 : tensor<1x2048xf32> loc(#loc5242)
    %3228 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3229 = stablehlo.reduce(%3227 init: %3228) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3230 = stablehlo.multiply %2813, %2813 : tensor<1x2048xf32> loc(#loc5242)
    %3231 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3232 = stablehlo.reduce(%3230 init: %3231) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3233 = stablehlo.multiply %2815, %2815 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3234 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3235 = stablehlo.reduce(%3233 init: %3234) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3236 = stablehlo.multiply %2817, %2817 : tensor<1x64xf32> loc(#loc5238)
    %3237 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3238 = stablehlo.reduce(%3236 init: %3237) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3239 = stablehlo.multiply %2819, %2819 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3240 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3241 = stablehlo.reduce(%3239 init: %3240) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3242 = stablehlo.multiply %2821, %2821 : tensor<1x8192xf32> loc(#loc5238)
    %3243 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3244 = stablehlo.reduce(%3242 init: %3243) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3245 = stablehlo.multiply %2823, %2823 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3246 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3247 = stablehlo.reduce(%3245 init: %3246) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3248 = stablehlo.multiply %2825, %2825 : tensor<1x2048xf32> loc(#loc5242)
    %3249 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3250 = stablehlo.reduce(%3248 init: %3249) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3251 = stablehlo.multiply %2827, %2827 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3252 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3253 = stablehlo.reduce(%3251 init: %3252) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3254 = stablehlo.multiply %2829, %2829 : tensor<1x2048xf32> loc(#loc5242)
    %3255 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3256 = stablehlo.reduce(%3254 init: %3255) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3257 = stablehlo.multiply %2831, %2831 : tensor<1x2048xf32> loc(#loc5242)
    %3258 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3259 = stablehlo.reduce(%3257 init: %3258) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3260 = stablehlo.multiply %2833, %2833 : tensor<1x2048xf32> loc(#loc5242)
    %3261 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3262 = stablehlo.reduce(%3260 init: %3261) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3263 = stablehlo.multiply %2835, %2835 : tensor<1x2048xf32> loc(#loc5242)
    %3264 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3265 = stablehlo.reduce(%3263 init: %3264) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3266 = stablehlo.multiply %2837, %2837 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3267 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3268 = stablehlo.reduce(%3266 init: %3267) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3269 = stablehlo.multiply %2839, %2839 : tensor<1x64xf32> loc(#loc5238)
    %3270 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3271 = stablehlo.reduce(%3269 init: %3270) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3272 = stablehlo.multiply %2841, %2841 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3273 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3274 = stablehlo.reduce(%3272 init: %3273) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3275 = stablehlo.multiply %2843, %2843 : tensor<1x8192xf32> loc(#loc5238)
    %3276 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3277 = stablehlo.reduce(%3275 init: %3276) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3278 = stablehlo.multiply %2845, %2845 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3279 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3280 = stablehlo.reduce(%3278 init: %3279) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3281 = stablehlo.multiply %2847, %2847 : tensor<1x2048xf32> loc(#loc5242)
    %3282 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3283 = stablehlo.reduce(%3281 init: %3282) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3284 = stablehlo.multiply %2849, %2849 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3285 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3286 = stablehlo.reduce(%3284 init: %3285) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3287 = stablehlo.multiply %2851, %2851 : tensor<1x2048xf32> loc(#loc5242)
    %3288 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3289 = stablehlo.reduce(%3287 init: %3288) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3290 = stablehlo.multiply %2853, %2853 : tensor<1x2048xf32> loc(#loc5242)
    %3291 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3292 = stablehlo.reduce(%3290 init: %3291) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3293 = stablehlo.multiply %2855, %2855 : tensor<1x2048xf32> loc(#loc5242)
    %3294 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3295 = stablehlo.reduce(%3293 init: %3294) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3296 = stablehlo.multiply %2857, %2857 : tensor<1x2048xf32> loc(#loc5242)
    %3297 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3298 = stablehlo.reduce(%3296 init: %3297) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3299 = stablehlo.multiply %2859, %2859 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3300 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3301 = stablehlo.reduce(%3299 init: %3300) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3302 = stablehlo.multiply %2861, %2861 : tensor<1x64xf32> loc(#loc5238)
    %3303 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3304 = stablehlo.reduce(%3302 init: %3303) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3305 = stablehlo.multiply %2863, %2863 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3306 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3307 = stablehlo.reduce(%3305 init: %3306) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3308 = stablehlo.multiply %2865, %2865 : tensor<1x8192xf32> loc(#loc5238)
    %3309 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3310 = stablehlo.reduce(%3308 init: %3309) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3311 = stablehlo.multiply %2867, %2867 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3312 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3313 = stablehlo.reduce(%3311 init: %3312) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3314 = stablehlo.multiply %2869, %2869 : tensor<1x2048xf32> loc(#loc5242)
    %3315 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3316 = stablehlo.reduce(%3314 init: %3315) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3317 = stablehlo.multiply %2871, %2871 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3318 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3319 = stablehlo.reduce(%3317 init: %3318) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3320 = stablehlo.multiply %2873, %2873 : tensor<1x2048xf32> loc(#loc5242)
    %3321 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3322 = stablehlo.reduce(%3320 init: %3321) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3323 = stablehlo.multiply %2875, %2875 : tensor<1x2048xf32> loc(#loc5242)
    %3324 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3325 = stablehlo.reduce(%3323 init: %3324) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3326 = stablehlo.multiply %2877, %2877 : tensor<1x2048xf32> loc(#loc5242)
    %3327 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3328 = stablehlo.reduce(%3326 init: %3327) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3329 = stablehlo.multiply %2879, %2879 : tensor<1x2048xf32> loc(#loc5242)
    %3330 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3331 = stablehlo.reduce(%3329 init: %3330) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3332 = stablehlo.multiply %2881, %2881 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3333 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3334 = stablehlo.reduce(%3332 init: %3333) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3335 = stablehlo.multiply %2883, %2883 : tensor<1x64xf32> loc(#loc5238)
    %3336 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3337 = stablehlo.reduce(%3335 init: %3336) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3338 = stablehlo.multiply %2885, %2885 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3339 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3340 = stablehlo.reduce(%3338 init: %3339) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3341 = stablehlo.multiply %2887, %2887 : tensor<1x8192xf32> loc(#loc5238)
    %3342 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3343 = stablehlo.reduce(%3341 init: %3342) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3344 = stablehlo.multiply %2889, %2889 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3345 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3346 = stablehlo.reduce(%3344 init: %3345) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3347 = stablehlo.multiply %2891, %2891 : tensor<1x2048xf32> loc(#loc5242)
    %3348 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3349 = stablehlo.reduce(%3347 init: %3348) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3350 = stablehlo.multiply %2893, %2893 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3351 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3352 = stablehlo.reduce(%3350 init: %3351) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3353 = stablehlo.multiply %2895, %2895 : tensor<1x2048xf32> loc(#loc5242)
    %3354 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3355 = stablehlo.reduce(%3353 init: %3354) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3356 = stablehlo.multiply %2897, %2897 : tensor<1x2048xf32> loc(#loc5242)
    %3357 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3358 = stablehlo.reduce(%3356 init: %3357) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3359 = stablehlo.multiply %2899, %2899 : tensor<1x2048xf32> loc(#loc5242)
    %3360 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3361 = stablehlo.reduce(%3359 init: %3360) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3362 = stablehlo.multiply %2901, %2901 : tensor<1x2048xf32> loc(#loc5242)
    %3363 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3364 = stablehlo.reduce(%3362 init: %3363) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3365 = stablehlo.multiply %2903, %2903 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3366 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3367 = stablehlo.reduce(%3365 init: %3366) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3368 = stablehlo.multiply %2905, %2905 : tensor<1x64xf32> loc(#loc5238)
    %3369 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3370 = stablehlo.reduce(%3368 init: %3369) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3371 = stablehlo.multiply %2907, %2907 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3372 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3373 = stablehlo.reduce(%3371 init: %3372) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3374 = stablehlo.multiply %2909, %2909 : tensor<1x8192xf32> loc(#loc5238)
    %3375 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3376 = stablehlo.reduce(%3374 init: %3375) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3377 = stablehlo.multiply %2911, %2911 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3378 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3379 = stablehlo.reduce(%3377 init: %3378) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3380 = stablehlo.multiply %2913, %2913 : tensor<1x2048xf32> loc(#loc5242)
    %3381 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3382 = stablehlo.reduce(%3380 init: %3381) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3383 = stablehlo.multiply %2915, %2915 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3384 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3385 = stablehlo.reduce(%3383 init: %3384) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3386 = stablehlo.multiply %2917, %2917 : tensor<1x2048xf32> loc(#loc5242)
    %3387 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3388 = stablehlo.reduce(%3386 init: %3387) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3389 = stablehlo.multiply %2919, %2919 : tensor<1x2048xf32> loc(#loc5242)
    %3390 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3391 = stablehlo.reduce(%3389 init: %3390) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3392 = stablehlo.multiply %2921, %2921 : tensor<1x2048xf32> loc(#loc5242)
    %3393 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3394 = stablehlo.reduce(%3392 init: %3393) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3395 = stablehlo.multiply %2923, %2923 : tensor<1x2048xf32> loc(#loc5242)
    %3396 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3397 = stablehlo.reduce(%3395 init: %3396) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3398 = stablehlo.multiply %2925, %2925 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3399 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3400 = stablehlo.reduce(%3398 init: %3399) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3401 = stablehlo.multiply %2927, %2927 : tensor<1x64xf32> loc(#loc5238)
    %3402 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3403 = stablehlo.reduce(%3401 init: %3402) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3404 = stablehlo.multiply %2929, %2929 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3405 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3406 = stablehlo.reduce(%3404 init: %3405) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3407 = stablehlo.multiply %2931, %2931 : tensor<1x8192xf32> loc(#loc5238)
    %3408 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3409 = stablehlo.reduce(%3407 init: %3408) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3410 = stablehlo.multiply %2933, %2933 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3411 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3412 = stablehlo.reduce(%3410 init: %3411) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3413 = stablehlo.multiply %2935, %2935 : tensor<1x2048xf32> loc(#loc5242)
    %3414 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3415 = stablehlo.reduce(%3413 init: %3414) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3416 = stablehlo.multiply %2937, %2937 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3417 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3418 = stablehlo.reduce(%3416 init: %3417) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3419 = stablehlo.multiply %2939, %2939 : tensor<1x2048xf32> loc(#loc5242)
    %3420 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3421 = stablehlo.reduce(%3419 init: %3420) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3422 = stablehlo.multiply %2941, %2941 : tensor<1x2048xf32> loc(#loc5242)
    %3423 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3424 = stablehlo.reduce(%3422 init: %3423) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3425 = stablehlo.multiply %2943, %2943 : tensor<1x2048xf32> loc(#loc5242)
    %3426 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3427 = stablehlo.reduce(%3425 init: %3426) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3428 = stablehlo.multiply %2945, %2945 : tensor<1x2048xf32> loc(#loc5242)
    %3429 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3430 = stablehlo.reduce(%3428 init: %3429) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3431 = stablehlo.multiply %2947, %2947 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3432 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3433 = stablehlo.reduce(%3431 init: %3432) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3434 = stablehlo.multiply %2949, %2949 : tensor<1x64xf32> loc(#loc5238)
    %3435 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3436 = stablehlo.reduce(%3434 init: %3435) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3437 = stablehlo.multiply %2951, %2951 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3438 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3439 = stablehlo.reduce(%3437 init: %3438) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3440 = stablehlo.multiply %2953, %2953 : tensor<1x8192xf32> loc(#loc5238)
    %3441 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3442 = stablehlo.reduce(%3440 init: %3441) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3443 = stablehlo.multiply %2955, %2955 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3444 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3445 = stablehlo.reduce(%3443 init: %3444) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3446 = stablehlo.multiply %2957, %2957 : tensor<1x2048xf32> loc(#loc5242)
    %3447 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3448 = stablehlo.reduce(%3446 init: %3447) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3449 = stablehlo.multiply %2959, %2959 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3450 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3451 = stablehlo.reduce(%3449 init: %3450) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3452 = stablehlo.multiply %2961, %2961 : tensor<1x2048xf32> loc(#loc5242)
    %3453 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3454 = stablehlo.reduce(%3452 init: %3453) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3455 = stablehlo.multiply %2963, %2963 : tensor<1x2048xf32> loc(#loc5242)
    %3456 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3457 = stablehlo.reduce(%3455 init: %3456) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3458 = stablehlo.multiply %2965, %2965 : tensor<1x2048xf32> loc(#loc5242)
    %3459 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3460 = stablehlo.reduce(%3458 init: %3459) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3461 = stablehlo.multiply %2967, %2967 : tensor<1x2048xf32> loc(#loc5242)
    %3462 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3463 = stablehlo.reduce(%3461 init: %3462) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3464 = stablehlo.multiply %2969, %2969 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3465 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3466 = stablehlo.reduce(%3464 init: %3465) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3467 = stablehlo.multiply %2971, %2971 : tensor<1x64xf32> loc(#loc5238)
    %3468 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3469 = stablehlo.reduce(%3467 init: %3468) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3470 = stablehlo.multiply %2973, %2973 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3471 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3472 = stablehlo.reduce(%3470 init: %3471) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3473 = stablehlo.multiply %2975, %2975 : tensor<1x8192xf32> loc(#loc5238)
    %3474 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3475 = stablehlo.reduce(%3473 init: %3474) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3476 = stablehlo.multiply %2977, %2977 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3477 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3478 = stablehlo.reduce(%3476 init: %3477) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3479 = stablehlo.multiply %2979, %2979 : tensor<1x2048xf32> loc(#loc5242)
    %3480 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3481 = stablehlo.reduce(%3479 init: %3480) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3482 = stablehlo.multiply %2981, %2981 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3483 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3484 = stablehlo.reduce(%3482 init: %3483) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3485 = stablehlo.multiply %2983, %2983 : tensor<1x2048xf32> loc(#loc5242)
    %3486 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3487 = stablehlo.reduce(%3485 init: %3486) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3488 = stablehlo.multiply %2985, %2985 : tensor<1x2048xf32> loc(#loc5242)
    %3489 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3490 = stablehlo.reduce(%3488 init: %3489) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3491 = stablehlo.multiply %2987, %2987 : tensor<1x2048xf32> loc(#loc5242)
    %3492 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3493 = stablehlo.reduce(%3491 init: %3492) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3494 = stablehlo.multiply %2989, %2989 : tensor<1x2048xf32> loc(#loc5242)
    %3495 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3496 = stablehlo.reduce(%3494 init: %3495) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3497 = stablehlo.multiply %2991, %2991 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3498 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3499 = stablehlo.reduce(%3497 init: %3498) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3500 = stablehlo.multiply %2993, %2993 : tensor<1x64xf32> loc(#loc5238)
    %3501 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3502 = stablehlo.reduce(%3500 init: %3501) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3503 = stablehlo.multiply %2995, %2995 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3504 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3505 = stablehlo.reduce(%3503 init: %3504) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3506 = stablehlo.multiply %2997, %2997 : tensor<1x8192xf32> loc(#loc5238)
    %3507 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3508 = stablehlo.reduce(%3506 init: %3507) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3509 = stablehlo.multiply %2999, %2999 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3510 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3511 = stablehlo.reduce(%3509 init: %3510) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3512 = stablehlo.multiply %3001, %3001 : tensor<1x2048xf32> loc(#loc5242)
    %3513 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3514 = stablehlo.reduce(%3512 init: %3513) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3515 = stablehlo.multiply %3003, %3003 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3516 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3517 = stablehlo.reduce(%3515 init: %3516) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3518 = stablehlo.multiply %3005, %3005 : tensor<1x2048xf32> loc(#loc5242)
    %3519 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3520 = stablehlo.reduce(%3518 init: %3519) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3521 = stablehlo.multiply %3007, %3007 : tensor<1x2048xf32> loc(#loc5242)
    %3522 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3523 = stablehlo.reduce(%3521 init: %3522) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3524 = stablehlo.multiply %3009, %3009 : tensor<1x2048xf32> loc(#loc5242)
    %3525 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3526 = stablehlo.reduce(%3524 init: %3525) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3527 = stablehlo.multiply %3011, %3011 : tensor<1x2048xf32> loc(#loc5242)
    %3528 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3529 = stablehlo.reduce(%3527 init: %3528) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3530 = stablehlo.multiply %3013, %3013 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3531 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3532 = stablehlo.reduce(%3530 init: %3531) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3533 = stablehlo.multiply %3015, %3015 : tensor<1x64xf32> loc(#loc5238)
    %3534 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3535 = stablehlo.reduce(%3533 init: %3534) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3536 = stablehlo.multiply %3017, %3017 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3537 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3538 = stablehlo.reduce(%3536 init: %3537) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3539 = stablehlo.multiply %3019, %3019 : tensor<1x8192xf32> loc(#loc5238)
    %3540 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3541 = stablehlo.reduce(%3539 init: %3540) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3542 = stablehlo.multiply %3021, %3021 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3543 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3544 = stablehlo.reduce(%3542 init: %3543) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3545 = stablehlo.multiply %3023, %3023 : tensor<1x2048xf32> loc(#loc5242)
    %3546 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3547 = stablehlo.reduce(%3545 init: %3546) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3548 = stablehlo.multiply %3025, %3025 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3549 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3550 = stablehlo.reduce(%3548 init: %3549) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3551 = stablehlo.multiply %3027, %3027 : tensor<1x2048xf32> loc(#loc5242)
    %3552 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3553 = stablehlo.reduce(%3551 init: %3552) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3554 = stablehlo.multiply %3029, %3029 : tensor<1x2048xf32> loc(#loc5242)
    %3555 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3556 = stablehlo.reduce(%3554 init: %3555) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3557 = stablehlo.multiply %3031, %3031 : tensor<1x2048xf32> loc(#loc5242)
    %3558 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3559 = stablehlo.reduce(%3557 init: %3558) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3560 = stablehlo.multiply %3033, %3033 : tensor<1x2048xf32> loc(#loc5242)
    %3561 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3562 = stablehlo.reduce(%3560 init: %3561) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3563 = stablehlo.multiply %3035, %3035 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3564 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3565 = stablehlo.reduce(%3563 init: %3564) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3566 = stablehlo.multiply %3037, %3037 : tensor<1x64xf32> loc(#loc5238)
    %3567 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3568 = stablehlo.reduce(%3566 init: %3567) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3569 = stablehlo.multiply %3039, %3039 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3570 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3571 = stablehlo.reduce(%3569 init: %3570) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3572 = stablehlo.multiply %3041, %3041 : tensor<1x8192xf32> loc(#loc5238)
    %3573 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3574 = stablehlo.reduce(%3572 init: %3573) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3575 = stablehlo.multiply %3043, %3043 : tensor<1x2048x8192xf32> loc(#loc5238)
    %3576 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3577 = stablehlo.reduce(%3575 init: %3576) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3578 = stablehlo.multiply %3045, %3045 : tensor<1x2048xf32> loc(#loc5242)
    %3579 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3580 = stablehlo.reduce(%3578 init: %3579) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3581 = stablehlo.multiply %3047, %3047 : tensor<1x8192x2048xf32> loc(#loc5238)
    %3582 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %3583 = stablehlo.reduce(%3581 init: %3582) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %3584 = stablehlo.multiply %3049, %3049 : tensor<1x2048xf32> loc(#loc5242)
    %3585 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3586 = stablehlo.reduce(%3584 init: %3585) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3587 = stablehlo.multiply %3051, %3051 : tensor<1x2048xf32> loc(#loc5242)
    %3588 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3589 = stablehlo.reduce(%3587 init: %3588) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3590 = stablehlo.multiply %3053, %3053 : tensor<1x2048xf32> loc(#loc5242)
    %3591 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3592 = stablehlo.reduce(%3590 init: %3591) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3593 = stablehlo.multiply %3055, %3055 : tensor<1x2048xf32> loc(#loc5242)
    %3594 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %3595 = stablehlo.reduce(%3593 init: %3594) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %3596 = stablehlo.multiply %3057, %3057 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %3597 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %3598 = stablehlo.reduce(%3596 init: %3597) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %3599 = stablehlo.multiply %3059, %3059 : tensor<1x64xf32> loc(#loc5238)
    %3600 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %3601 = stablehlo.reduce(%3599 init: %3600) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %3602 = stablehlo.multiply %3061, %3061 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %3603 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %3604 = stablehlo.reduce(%3602 init: %3603) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %3605 = stablehlo.broadcast_in_dim %3064, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3606 = stablehlo.broadcast_in_dim %3067, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3607 = stablehlo.broadcast_in_dim %3070, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3608 = stablehlo.broadcast_in_dim %3073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3609 = stablehlo.broadcast_in_dim %3076, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3610 = stablehlo.broadcast_in_dim %3079, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3611 = stablehlo.broadcast_in_dim %3082, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3612 = stablehlo.broadcast_in_dim %3085, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3613 = stablehlo.broadcast_in_dim %3088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3614 = stablehlo.broadcast_in_dim %3091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3615 = stablehlo.broadcast_in_dim %3094, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3616 = stablehlo.broadcast_in_dim %3097, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3617 = stablehlo.broadcast_in_dim %3100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3618 = stablehlo.broadcast_in_dim %3103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3619 = stablehlo.broadcast_in_dim %3106, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3620 = stablehlo.broadcast_in_dim %3109, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3621 = stablehlo.broadcast_in_dim %3112, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3622 = stablehlo.broadcast_in_dim %3115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3623 = stablehlo.broadcast_in_dim %3118, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3624 = stablehlo.broadcast_in_dim %3121, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3625 = stablehlo.broadcast_in_dim %3124, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3626 = stablehlo.broadcast_in_dim %3127, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3627 = stablehlo.broadcast_in_dim %3130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3628 = stablehlo.broadcast_in_dim %3133, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3629 = stablehlo.broadcast_in_dim %3136, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3630 = stablehlo.broadcast_in_dim %3139, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3631 = stablehlo.broadcast_in_dim %3142, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3632 = stablehlo.broadcast_in_dim %3145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3633 = stablehlo.broadcast_in_dim %3148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3634 = stablehlo.broadcast_in_dim %3151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3635 = stablehlo.broadcast_in_dim %3154, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3636 = stablehlo.broadcast_in_dim %3157, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3637 = stablehlo.broadcast_in_dim %3160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3638 = stablehlo.broadcast_in_dim %3163, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3639 = stablehlo.broadcast_in_dim %3166, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3640 = stablehlo.broadcast_in_dim %3169, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3641 = stablehlo.broadcast_in_dim %3172, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3642 = stablehlo.broadcast_in_dim %3175, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3643 = stablehlo.broadcast_in_dim %3178, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3644 = stablehlo.broadcast_in_dim %3181, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3645 = stablehlo.broadcast_in_dim %3184, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3646 = stablehlo.broadcast_in_dim %3187, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3647 = stablehlo.broadcast_in_dim %3190, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3648 = stablehlo.broadcast_in_dim %3193, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3649 = stablehlo.broadcast_in_dim %3196, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3650 = stablehlo.broadcast_in_dim %3199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3651 = stablehlo.broadcast_in_dim %3202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3652 = stablehlo.broadcast_in_dim %3205, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3653 = stablehlo.broadcast_in_dim %3208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3654 = stablehlo.broadcast_in_dim %3211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3655 = stablehlo.broadcast_in_dim %3214, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3656 = stablehlo.broadcast_in_dim %3217, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3657 = stablehlo.broadcast_in_dim %3220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3658 = stablehlo.broadcast_in_dim %3223, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3659 = stablehlo.broadcast_in_dim %3226, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3660 = stablehlo.broadcast_in_dim %3229, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3661 = stablehlo.broadcast_in_dim %3232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3662 = stablehlo.broadcast_in_dim %3235, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3663 = stablehlo.broadcast_in_dim %3238, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3664 = stablehlo.broadcast_in_dim %3241, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3665 = stablehlo.broadcast_in_dim %3244, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3666 = stablehlo.broadcast_in_dim %3247, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3667 = stablehlo.broadcast_in_dim %3250, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3668 = stablehlo.broadcast_in_dim %3253, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3669 = stablehlo.broadcast_in_dim %3256, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3670 = stablehlo.broadcast_in_dim %3259, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3671 = stablehlo.broadcast_in_dim %3262, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3672 = stablehlo.broadcast_in_dim %3265, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3673 = stablehlo.broadcast_in_dim %3268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3674 = stablehlo.broadcast_in_dim %3271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3675 = stablehlo.broadcast_in_dim %3274, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3676 = stablehlo.broadcast_in_dim %3277, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3677 = stablehlo.broadcast_in_dim %3280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3678 = stablehlo.broadcast_in_dim %3283, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3679 = stablehlo.broadcast_in_dim %3286, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3680 = stablehlo.broadcast_in_dim %3289, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3681 = stablehlo.broadcast_in_dim %3292, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3682 = stablehlo.broadcast_in_dim %3295, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3683 = stablehlo.broadcast_in_dim %3298, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3684 = stablehlo.broadcast_in_dim %3301, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3685 = stablehlo.broadcast_in_dim %3304, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3686 = stablehlo.broadcast_in_dim %3307, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3687 = stablehlo.broadcast_in_dim %3310, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3688 = stablehlo.broadcast_in_dim %3313, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3689 = stablehlo.broadcast_in_dim %3316, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3690 = stablehlo.broadcast_in_dim %3319, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3691 = stablehlo.broadcast_in_dim %3322, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3692 = stablehlo.broadcast_in_dim %3325, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3693 = stablehlo.broadcast_in_dim %3328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3694 = stablehlo.broadcast_in_dim %3331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3695 = stablehlo.broadcast_in_dim %3334, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3696 = stablehlo.broadcast_in_dim %3337, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3697 = stablehlo.broadcast_in_dim %3340, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3698 = stablehlo.broadcast_in_dim %3343, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3699 = stablehlo.broadcast_in_dim %3346, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3700 = stablehlo.broadcast_in_dim %3349, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3701 = stablehlo.broadcast_in_dim %3352, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3702 = stablehlo.broadcast_in_dim %3355, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3703 = stablehlo.broadcast_in_dim %3358, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3704 = stablehlo.broadcast_in_dim %3361, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3705 = stablehlo.broadcast_in_dim %3364, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3706 = stablehlo.broadcast_in_dim %3367, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3707 = stablehlo.broadcast_in_dim %3370, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3708 = stablehlo.broadcast_in_dim %3373, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3709 = stablehlo.broadcast_in_dim %3376, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3710 = stablehlo.broadcast_in_dim %3379, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3711 = stablehlo.broadcast_in_dim %3382, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3712 = stablehlo.broadcast_in_dim %3385, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3713 = stablehlo.broadcast_in_dim %3388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3714 = stablehlo.broadcast_in_dim %3391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3715 = stablehlo.broadcast_in_dim %3394, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3716 = stablehlo.broadcast_in_dim %3397, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3717 = stablehlo.broadcast_in_dim %3400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3718 = stablehlo.broadcast_in_dim %3403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3719 = stablehlo.broadcast_in_dim %3406, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3720 = stablehlo.broadcast_in_dim %3409, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3721 = stablehlo.broadcast_in_dim %3412, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3722 = stablehlo.broadcast_in_dim %3415, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3723 = stablehlo.broadcast_in_dim %3418, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3724 = stablehlo.broadcast_in_dim %3421, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3725 = stablehlo.broadcast_in_dim %3424, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3726 = stablehlo.broadcast_in_dim %3427, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3727 = stablehlo.broadcast_in_dim %3430, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3728 = stablehlo.broadcast_in_dim %3433, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3729 = stablehlo.broadcast_in_dim %3436, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3730 = stablehlo.broadcast_in_dim %3439, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3731 = stablehlo.broadcast_in_dim %3442, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3732 = stablehlo.broadcast_in_dim %3445, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3733 = stablehlo.broadcast_in_dim %3448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3734 = stablehlo.broadcast_in_dim %3451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3735 = stablehlo.broadcast_in_dim %3454, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3736 = stablehlo.broadcast_in_dim %3457, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3737 = stablehlo.broadcast_in_dim %3460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3738 = stablehlo.broadcast_in_dim %3463, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3739 = stablehlo.broadcast_in_dim %3466, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3740 = stablehlo.broadcast_in_dim %3469, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3741 = stablehlo.broadcast_in_dim %3472, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3742 = stablehlo.broadcast_in_dim %3475, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3743 = stablehlo.broadcast_in_dim %3478, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3744 = stablehlo.broadcast_in_dim %3481, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3745 = stablehlo.broadcast_in_dim %3484, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3746 = stablehlo.broadcast_in_dim %3487, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3747 = stablehlo.broadcast_in_dim %3490, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3748 = stablehlo.broadcast_in_dim %3493, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3749 = stablehlo.broadcast_in_dim %3496, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3750 = stablehlo.broadcast_in_dim %3499, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3751 = stablehlo.broadcast_in_dim %3502, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3752 = stablehlo.broadcast_in_dim %3505, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3753 = stablehlo.broadcast_in_dim %3508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3754 = stablehlo.broadcast_in_dim %3511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3755 = stablehlo.broadcast_in_dim %3514, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3756 = stablehlo.broadcast_in_dim %3517, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3757 = stablehlo.broadcast_in_dim %3520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3758 = stablehlo.broadcast_in_dim %3523, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3759 = stablehlo.broadcast_in_dim %3526, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3760 = stablehlo.broadcast_in_dim %3529, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3761 = stablehlo.broadcast_in_dim %3532, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3762 = stablehlo.broadcast_in_dim %3535, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3763 = stablehlo.broadcast_in_dim %3538, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3764 = stablehlo.broadcast_in_dim %3541, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3765 = stablehlo.broadcast_in_dim %3544, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3766 = stablehlo.broadcast_in_dim %3547, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3767 = stablehlo.broadcast_in_dim %3550, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3768 = stablehlo.broadcast_in_dim %3553, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3769 = stablehlo.broadcast_in_dim %3556, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3770 = stablehlo.broadcast_in_dim %3559, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3771 = stablehlo.broadcast_in_dim %3562, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3772 = stablehlo.broadcast_in_dim %3565, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3773 = stablehlo.broadcast_in_dim %3568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3774 = stablehlo.broadcast_in_dim %3571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3775 = stablehlo.broadcast_in_dim %3574, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3776 = stablehlo.broadcast_in_dim %3577, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3777 = stablehlo.broadcast_in_dim %3580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3778 = stablehlo.broadcast_in_dim %3583, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3779 = stablehlo.broadcast_in_dim %3586, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3780 = stablehlo.broadcast_in_dim %3589, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3781 = stablehlo.broadcast_in_dim %3592, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3782 = stablehlo.broadcast_in_dim %3595, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3783 = stablehlo.broadcast_in_dim %3598, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3784 = stablehlo.broadcast_in_dim %3601, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3785 = stablehlo.broadcast_in_dim %3604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %3786 = stablehlo.concatenate %3605, %3606, %3607, %3608, %3609, %3610, %3611, %3612, %3613, %3614, %3615, %3616, %3617, %3618, %3619, %3620, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3787 = stablehlo.concatenate %3621, %3622, %3623, %3624, %3625, %3626, %3627, %3628, %3629, %3630, %3631, %3632, %3633, %3634, %3635, %3636, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3788 = stablehlo.concatenate %3637, %3638, %3639, %3640, %3641, %3642, %3643, %3644, %3645, %3646, %3647, %3648, %3649, %3650, %3651, %3652, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3789 = stablehlo.concatenate %3653, %3654, %3655, %3656, %3657, %3658, %3659, %3660, %3661, %3662, %3663, %3664, %3665, %3666, %3667, %3668, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3790 = stablehlo.concatenate %3669, %3670, %3671, %3672, %3673, %3674, %3675, %3676, %3677, %3678, %3679, %3680, %3681, %3682, %3683, %3684, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3791 = stablehlo.concatenate %3685, %3686, %3687, %3688, %3689, %3690, %3691, %3692, %3693, %3694, %3695, %3696, %3697, %3698, %3699, %3700, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3792 = stablehlo.concatenate %3701, %3702, %3703, %3704, %3705, %3706, %3707, %3708, %3709, %3710, %3711, %3712, %3713, %3714, %3715, %3716, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3793 = stablehlo.concatenate %3717, %3718, %3719, %3720, %3721, %3722, %3723, %3724, %3725, %3726, %3727, %3728, %3729, %3730, %3731, %3732, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3794 = stablehlo.concatenate %3733, %3734, %3735, %3736, %3737, %3738, %3739, %3740, %3741, %3742, %3743, %3744, %3745, %3746, %3747, %3748, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3795 = stablehlo.concatenate %3749, %3750, %3751, %3752, %3753, %3754, %3755, %3756, %3757, %3758, %3759, %3760, %3761, %3762, %3763, %3764, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3796 = stablehlo.concatenate %3765, %3766, %3767, %3768, %3769, %3770, %3771, %3772, %3773, %3774, %3775, %3776, %3777, %3778, %3779, %3780, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %3797 = stablehlo.concatenate %3781, %3782, %3783, %3784, %3785, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5xf32> loc(#loc5247)
    %3798 = stablehlo.concatenate %3786, %3787, %3788, %3789, %3790, %3791, %3792, %3793, %3794, %3795, %3796, %3797, dim = 0 : (tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<5xf32>) -> tensor<181xf32> loc(#loc5247)
    %3799 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5248)
    %3800 = stablehlo.reduce(%3798 init: %3799) across dimensions = [0] : (tensor<181xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5248)
      stablehlo.return %29929 : tensor<f32> loc(#loc5248)
    } loc(#loc5248)
    %3801 = stablehlo.sqrt %3800 : tensor<f32> loc(#loc5249)
    %3802 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %3803 = stablehlo.add %arg182, %3802 : tensor<i32> loc(#loc5256)
    %3804 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %3805 = stablehlo.add %arg183, %3804 : tensor<i32> loc(#loc5256)
    %3806 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3807 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3808 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3809 = call @nan_to_num(%2701, %3806, %3807, %3808) : (tensor<2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<2048xf32> loc(#loc5257)
    %3810 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3811 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3812 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3813 = call @nan_to_num_1148(%2703, %3810, %3811, %3812) : (tensor<2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<2048xf32> loc(#loc5257)
    %3814 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3815 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3816 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3817 = call @nan_to_num_1153(%2705, %3814, %3815, %3816) : (tensor<2048x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5257)
    %3818 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3819 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3820 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3821 = call @nan_to_num_1158(%2707, %3818, %3819, %3820) : (tensor<51200xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<51200xf32> loc(#loc5257)
    %3822 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3823 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3824 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %3825 = call @nan_to_num_1163(%2709, %3822, %3823, %3824) : (tensor<2048x51200xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5257)
    %3826 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3827 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3828 = stablehlo.add %3826, %3827 : tensor<f32> loc(#loc5259)
    %3829 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3830 = stablehlo.subtract %3828, %3829 : tensor<f32> loc(#loc5260)
    %3831 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3832 = stablehlo.power %3831, %3830 : tensor<f32> loc(#loc5261)
    %3833 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3834 = stablehlo.subtract %3833, %3832 : tensor<f32> loc(#loc5260)
    %3835 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3836 = stablehlo.multiply %3835, %3834 : tensor<f32> loc(#loc5236)
    %3837 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3838 = stablehlo.power %3837, %3828 : tensor<f32> loc(#loc5261)
    %3839 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3840 = stablehlo.subtract %3839, %3838 : tensor<f32> loc(#loc5260)
    %3841 = stablehlo.divide %3836, %3840 : tensor<f32> loc(#loc5253)
    %3842 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3843 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3844 = stablehlo.add %3842, %3843 : tensor<f32> loc(#loc5259)
    %3845 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3846 = stablehlo.subtract %3844, %3845 : tensor<f32> loc(#loc5260)
    %3847 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3848 = stablehlo.power %3847, %3846 : tensor<f32> loc(#loc5261)
    %3849 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3850 = stablehlo.subtract %3849, %3848 : tensor<f32> loc(#loc5260)
    %3851 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3852 = stablehlo.multiply %3851, %3850 : tensor<f32> loc(#loc5236)
    %3853 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3854 = stablehlo.power %3853, %3844 : tensor<f32> loc(#loc5261)
    %3855 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3856 = stablehlo.subtract %3855, %3854 : tensor<f32> loc(#loc5260)
    %3857 = stablehlo.divide %3852, %3856 : tensor<f32> loc(#loc5253)
    %3858 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3859 = stablehlo.subtract %3858, %3841 : tensor<f32> loc(#loc5260)
    %3860 = stablehlo.broadcast_in_dim %3859, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3861 = stablehlo.multiply %3860, %3809 : tensor<2048xf32> loc(#loc5262)
    %3862 = stablehlo.broadcast_in_dim %3841, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3863 = stablehlo.multiply %3862, %arg185 : tensor<2048xf32> loc(#loc5262)
    %3864 = stablehlo.add %3861, %3863 : tensor<2048xf32> loc(#loc5263)
    %3865 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3866 = stablehlo.subtract %3865, %3857 : tensor<f32> loc(#loc5260)
    %3867 = stablehlo.multiply %3809, %3809 : tensor<2048xf32> loc(#loc5264)
    %3868 = stablehlo.broadcast_in_dim %3866, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3869 = stablehlo.multiply %3868, %3867 : tensor<2048xf32> loc(#loc5262)
    %3870 = stablehlo.broadcast_in_dim %3857, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3871 = stablehlo.multiply %3870, %arg190 : tensor<2048xf32> loc(#loc5262)
    %3872 = stablehlo.add %3869, %3871 : tensor<2048xf32> loc(#loc5263)
    %3873 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3874 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3875 = stablehlo.add %3873, %3874 : tensor<f32> loc(#loc5259)
    %3876 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3877 = stablehlo.subtract %3875, %3876 : tensor<f32> loc(#loc5260)
    %3878 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3879 = stablehlo.power %3878, %3877 : tensor<f32> loc(#loc5261)
    %3880 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3881 = stablehlo.subtract %3880, %3879 : tensor<f32> loc(#loc5260)
    %3882 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3883 = stablehlo.multiply %3882, %3881 : tensor<f32> loc(#loc5236)
    %3884 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3885 = stablehlo.power %3884, %3875 : tensor<f32> loc(#loc5261)
    %3886 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3887 = stablehlo.subtract %3886, %3885 : tensor<f32> loc(#loc5260)
    %3888 = stablehlo.divide %3883, %3887 : tensor<f32> loc(#loc5253)
    %3889 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3890 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3891 = stablehlo.add %3889, %3890 : tensor<f32> loc(#loc5259)
    %3892 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3893 = stablehlo.subtract %3891, %3892 : tensor<f32> loc(#loc5260)
    %3894 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3895 = stablehlo.power %3894, %3893 : tensor<f32> loc(#loc5261)
    %3896 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3897 = stablehlo.subtract %3896, %3895 : tensor<f32> loc(#loc5260)
    %3898 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3899 = stablehlo.multiply %3898, %3897 : tensor<f32> loc(#loc5236)
    %3900 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3901 = stablehlo.power %3900, %3891 : tensor<f32> loc(#loc5261)
    %3902 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3903 = stablehlo.subtract %3902, %3901 : tensor<f32> loc(#loc5260)
    %3904 = stablehlo.divide %3899, %3903 : tensor<f32> loc(#loc5253)
    %3905 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3906 = stablehlo.subtract %3905, %3888 : tensor<f32> loc(#loc5260)
    %3907 = stablehlo.broadcast_in_dim %3906, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3908 = stablehlo.multiply %3907, %3813 : tensor<2048xf32> loc(#loc5262)
    %3909 = stablehlo.broadcast_in_dim %3888, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3910 = stablehlo.multiply %3909, %arg186 : tensor<2048xf32> loc(#loc5262)
    %3911 = stablehlo.add %3908, %3910 : tensor<2048xf32> loc(#loc5263)
    %3912 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3913 = stablehlo.subtract %3912, %3904 : tensor<f32> loc(#loc5260)
    %3914 = stablehlo.multiply %3813, %3813 : tensor<2048xf32> loc(#loc5264)
    %3915 = stablehlo.broadcast_in_dim %3913, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3916 = stablehlo.multiply %3915, %3914 : tensor<2048xf32> loc(#loc5262)
    %3917 = stablehlo.broadcast_in_dim %3904, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %3918 = stablehlo.multiply %3917, %arg191 : tensor<2048xf32> loc(#loc5262)
    %3919 = stablehlo.add %3916, %3918 : tensor<2048xf32> loc(#loc5263)
    %3920 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3921 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3922 = stablehlo.add %3920, %3921 : tensor<f32> loc(#loc5259)
    %3923 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3924 = stablehlo.subtract %3922, %3923 : tensor<f32> loc(#loc5260)
    %3925 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3926 = stablehlo.power %3925, %3924 : tensor<f32> loc(#loc5261)
    %3927 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3928 = stablehlo.subtract %3927, %3926 : tensor<f32> loc(#loc5260)
    %3929 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3930 = stablehlo.multiply %3929, %3928 : tensor<f32> loc(#loc5236)
    %3931 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3932 = stablehlo.power %3931, %3922 : tensor<f32> loc(#loc5261)
    %3933 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3934 = stablehlo.subtract %3933, %3932 : tensor<f32> loc(#loc5260)
    %3935 = stablehlo.divide %3930, %3934 : tensor<f32> loc(#loc5253)
    %3936 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3937 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3938 = stablehlo.add %3936, %3937 : tensor<f32> loc(#loc5259)
    %3939 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3940 = stablehlo.subtract %3938, %3939 : tensor<f32> loc(#loc5260)
    %3941 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3942 = stablehlo.power %3941, %3940 : tensor<f32> loc(#loc5261)
    %3943 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3944 = stablehlo.subtract %3943, %3942 : tensor<f32> loc(#loc5260)
    %3945 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3946 = stablehlo.multiply %3945, %3944 : tensor<f32> loc(#loc5236)
    %3947 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3948 = stablehlo.power %3947, %3938 : tensor<f32> loc(#loc5261)
    %3949 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3950 = stablehlo.subtract %3949, %3948 : tensor<f32> loc(#loc5260)
    %3951 = stablehlo.divide %3946, %3950 : tensor<f32> loc(#loc5253)
    %3952 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3953 = stablehlo.subtract %3952, %3935 : tensor<f32> loc(#loc5260)
    %3954 = stablehlo.broadcast_in_dim %3953, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5262)
    %3955 = stablehlo.multiply %3954, %3817 : tensor<2048x2048xf32> loc(#loc5262)
    %3956 = stablehlo.broadcast_in_dim %3935, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5262)
    %3957 = stablehlo.multiply %3956, %arg187 : tensor<2048x2048xf32> loc(#loc5262)
    %3958 = stablehlo.add %3955, %3957 : tensor<2048x2048xf32> loc(#loc5263)
    %3959 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3960 = stablehlo.subtract %3959, %3951 : tensor<f32> loc(#loc5260)
    %3961 = stablehlo.multiply %3817, %3817 : tensor<2048x2048xf32> loc(#loc5264)
    %3962 = stablehlo.broadcast_in_dim %3960, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5262)
    %3963 = stablehlo.multiply %3962, %3961 : tensor<2048x2048xf32> loc(#loc5262)
    %3964 = stablehlo.broadcast_in_dim %3951, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5262)
    %3965 = stablehlo.multiply %3964, %arg192 : tensor<2048x2048xf32> loc(#loc5262)
    %3966 = stablehlo.add %3963, %3965 : tensor<2048x2048xf32> loc(#loc5263)
    %3967 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3969 = stablehlo.add %3967, %3968 : tensor<f32> loc(#loc5259)
    %3970 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3971 = stablehlo.subtract %3969, %3970 : tensor<f32> loc(#loc5260)
    %3972 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3973 = stablehlo.power %3972, %3971 : tensor<f32> loc(#loc5261)
    %3974 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3975 = stablehlo.subtract %3974, %3973 : tensor<f32> loc(#loc5260)
    %3976 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3977 = stablehlo.multiply %3976, %3975 : tensor<f32> loc(#loc5236)
    %3978 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %3979 = stablehlo.power %3978, %3969 : tensor<f32> loc(#loc5261)
    %3980 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3981 = stablehlo.subtract %3980, %3979 : tensor<f32> loc(#loc5260)
    %3982 = stablehlo.divide %3977, %3981 : tensor<f32> loc(#loc5253)
    %3983 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %3984 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3985 = stablehlo.add %3983, %3984 : tensor<f32> loc(#loc5259)
    %3986 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3987 = stablehlo.subtract %3985, %3986 : tensor<f32> loc(#loc5260)
    %3988 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3989 = stablehlo.power %3988, %3987 : tensor<f32> loc(#loc5261)
    %3990 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3991 = stablehlo.subtract %3990, %3989 : tensor<f32> loc(#loc5260)
    %3992 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3993 = stablehlo.multiply %3992, %3991 : tensor<f32> loc(#loc5236)
    %3994 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %3995 = stablehlo.power %3994, %3985 : tensor<f32> loc(#loc5261)
    %3996 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %3997 = stablehlo.subtract %3996, %3995 : tensor<f32> loc(#loc5260)
    %3998 = stablehlo.divide %3993, %3997 : tensor<f32> loc(#loc5253)
    %3999 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4000 = stablehlo.subtract %3999, %3982 : tensor<f32> loc(#loc5260)
    %4001 = stablehlo.broadcast_in_dim %4000, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5262)
    %4002 = stablehlo.multiply %4001, %3821 : tensor<51200xf32> loc(#loc5262)
    %4003 = stablehlo.broadcast_in_dim %3982, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5262)
    %4004 = stablehlo.multiply %4003, %arg188 : tensor<51200xf32> loc(#loc5262)
    %4005 = stablehlo.add %4002, %4004 : tensor<51200xf32> loc(#loc5263)
    %4006 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4007 = stablehlo.subtract %4006, %3998 : tensor<f32> loc(#loc5260)
    %4008 = stablehlo.multiply %3821, %3821 : tensor<51200xf32> loc(#loc5264)
    %4009 = stablehlo.broadcast_in_dim %4007, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5262)
    %4010 = stablehlo.multiply %4009, %4008 : tensor<51200xf32> loc(#loc5262)
    %4011 = stablehlo.broadcast_in_dim %3998, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5262)
    %4012 = stablehlo.multiply %4011, %arg193 : tensor<51200xf32> loc(#loc5262)
    %4013 = stablehlo.add %4010, %4012 : tensor<51200xf32> loc(#loc5263)
    %4014 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %4015 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4016 = stablehlo.add %4014, %4015 : tensor<f32> loc(#loc5259)
    %4017 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4018 = stablehlo.subtract %4016, %4017 : tensor<f32> loc(#loc5260)
    %4019 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %4020 = stablehlo.power %4019, %4018 : tensor<f32> loc(#loc5261)
    %4021 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4022 = stablehlo.subtract %4021, %4020 : tensor<f32> loc(#loc5260)
    %4023 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %4024 = stablehlo.multiply %4023, %4022 : tensor<f32> loc(#loc5236)
    %4025 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %4026 = stablehlo.power %4025, %4016 : tensor<f32> loc(#loc5261)
    %4027 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4028 = stablehlo.subtract %4027, %4026 : tensor<f32> loc(#loc5260)
    %4029 = stablehlo.divide %4024, %4028 : tensor<f32> loc(#loc5253)
    %4030 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5258)
    %4031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4032 = stablehlo.add %4030, %4031 : tensor<f32> loc(#loc5259)
    %4033 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4034 = stablehlo.subtract %4032, %4033 : tensor<f32> loc(#loc5260)
    %4035 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %4036 = stablehlo.power %4035, %4034 : tensor<f32> loc(#loc5261)
    %4037 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4038 = stablehlo.subtract %4037, %4036 : tensor<f32> loc(#loc5260)
    %4039 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %4040 = stablehlo.multiply %4039, %4038 : tensor<f32> loc(#loc5236)
    %4041 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %4042 = stablehlo.power %4041, %4032 : tensor<f32> loc(#loc5261)
    %4043 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4044 = stablehlo.subtract %4043, %4042 : tensor<f32> loc(#loc5260)
    %4045 = stablehlo.divide %4040, %4044 : tensor<f32> loc(#loc5253)
    %4046 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4047 = stablehlo.subtract %4046, %4029 : tensor<f32> loc(#loc5260)
    %4048 = stablehlo.broadcast_in_dim %4047, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5262)
    %4049 = stablehlo.multiply %4048, %3825 : tensor<2048x51200xf32> loc(#loc5262)
    %4050 = stablehlo.broadcast_in_dim %4029, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5262)
    %4051 = stablehlo.multiply %4050, %arg189 : tensor<2048x51200xf32> loc(#loc5262)
    %4052 = stablehlo.add %4049, %4051 : tensor<2048x51200xf32> loc(#loc5263)
    %4053 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4054 = stablehlo.subtract %4053, %4045 : tensor<f32> loc(#loc5260)
    %4055 = stablehlo.multiply %3825, %3825 : tensor<2048x51200xf32> loc(#loc5264)
    %4056 = stablehlo.broadcast_in_dim %4054, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5262)
    %4057 = stablehlo.multiply %4056, %4055 : tensor<2048x51200xf32> loc(#loc5262)
    %4058 = stablehlo.broadcast_in_dim %4045, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5262)
    %4059 = stablehlo.multiply %4058, %arg194 : tensor<2048x51200xf32> loc(#loc5262)
    %4060 = stablehlo.add %4057, %4059 : tensor<2048x51200xf32> loc(#loc5263)
    %4061 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4062 = stablehlo.broadcast_in_dim %4061, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5265)
    %4063 = stablehlo.add %3872, %4062 : tensor<2048xf32> loc(#loc5265)
    %4064 = stablehlo.sqrt %4063 : tensor<2048xf32> loc(#loc5266)
    %4065 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %4066 = stablehlo.broadcast_in_dim %4065, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5265)
    %4067 = stablehlo.add %4064, %4066 : tensor<2048xf32> loc(#loc5265)
    %4068 = stablehlo.divide %3864, %4067 : tensor<2048xf32> loc(#loc5267)
    %4069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4070 = stablehlo.broadcast_in_dim %4069, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5265)
    %4071 = stablehlo.add %3919, %4070 : tensor<2048xf32> loc(#loc5265)
    %4072 = stablehlo.sqrt %4071 : tensor<2048xf32> loc(#loc5266)
    %4073 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %4074 = stablehlo.broadcast_in_dim %4073, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5265)
    %4075 = stablehlo.add %4072, %4074 : tensor<2048xf32> loc(#loc5265)
    %4076 = stablehlo.divide %3911, %4075 : tensor<2048xf32> loc(#loc5267)
    %4077 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4078 = stablehlo.broadcast_in_dim %4077, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5265)
    %4079 = stablehlo.add %3966, %4078 : tensor<2048x2048xf32> loc(#loc5265)
    %4080 = stablehlo.sqrt %4079 : tensor<2048x2048xf32> loc(#loc5266)
    %4081 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %4082 = stablehlo.broadcast_in_dim %4081, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5265)
    %4083 = stablehlo.add %4080, %4082 : tensor<2048x2048xf32> loc(#loc5265)
    %4084 = stablehlo.divide %3958, %4083 : tensor<2048x2048xf32> loc(#loc5267)
    %4085 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4086 = stablehlo.broadcast_in_dim %4085, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5265)
    %4087 = stablehlo.add %4013, %4086 : tensor<51200xf32> loc(#loc5265)
    %4088 = stablehlo.sqrt %4087 : tensor<51200xf32> loc(#loc5266)
    %4089 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %4090 = stablehlo.broadcast_in_dim %4089, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5265)
    %4091 = stablehlo.add %4088, %4090 : tensor<51200xf32> loc(#loc5265)
    %4092 = stablehlo.divide %4005, %4091 : tensor<51200xf32> loc(#loc5267)
    %4093 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4094 = stablehlo.broadcast_in_dim %4093, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5265)
    %4095 = stablehlo.add %4060, %4094 : tensor<2048x51200xf32> loc(#loc5265)
    %4096 = stablehlo.sqrt %4095 : tensor<2048x51200xf32> loc(#loc5266)
    %4097 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %4098 = stablehlo.broadcast_in_dim %4097, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5265)
    %4099 = stablehlo.add %4096, %4098 : tensor<2048x51200xf32> loc(#loc5265)
    %4100 = stablehlo.divide %4052, %4099 : tensor<2048x51200xf32> loc(#loc5267)
    %4101 = stablehlo.multiply %4068, %4068 : tensor<2048xf32> loc(#loc5268)
    %4102 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5269)
    %4103 = stablehlo.reduce(%4101 init: %4102) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5269)
      stablehlo.return %29929 : tensor<f32> loc(#loc5269)
    } loc(#loc5269)
    %4104 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %4105 = stablehlo.divide %4103, %4104 : tensor<f32> loc(#loc5270)
    %4106 = stablehlo.sqrt %4105 : tensor<f32> loc(#loc5249)
    %4107 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4108 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4109 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4110 = call @nan_to_num_1168(%4106, %4107, %4108, %4109) : (tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5257)
    %4111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4112 = stablehlo.divide %4110, %4111 : tensor<f32> loc(#loc5253)
    %4113 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4114 = stablehlo.maximum %4113, %4112 : tensor<f32> loc(#loc5271)
    %4115 = stablehlo.broadcast_in_dim %4114, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5272)
    %4116 = stablehlo.divide %4068, %4115 : tensor<2048xf32> loc(#loc5272)
    %4117 = stablehlo.multiply %4076, %4076 : tensor<2048xf32> loc(#loc5268)
    %4118 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5269)
    %4119 = stablehlo.reduce(%4117 init: %4118) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5269)
      stablehlo.return %29929 : tensor<f32> loc(#loc5269)
    } loc(#loc5269)
    %4120 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %4121 = stablehlo.divide %4119, %4120 : tensor<f32> loc(#loc5270)
    %4122 = stablehlo.sqrt %4121 : tensor<f32> loc(#loc5249)
    %4123 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4124 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4125 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4126 = call @nan_to_num_1173(%4122, %4123, %4124, %4125) : (tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5257)
    %4127 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4128 = stablehlo.divide %4126, %4127 : tensor<f32> loc(#loc5253)
    %4129 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4130 = stablehlo.maximum %4129, %4128 : tensor<f32> loc(#loc5271)
    %4131 = stablehlo.broadcast_in_dim %4130, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5272)
    %4132 = stablehlo.divide %4076, %4131 : tensor<2048xf32> loc(#loc5272)
    %4133 = stablehlo.multiply %4084, %4084 : tensor<2048x2048xf32> loc(#loc5268)
    %4134 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5273)
    %4135 = stablehlo.reduce(%4133 init: %4134) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5273)
      stablehlo.return %29929 : tensor<f32> loc(#loc5273)
    } loc(#loc5273)
    %4136 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %4137 = stablehlo.divide %4135, %4136 : tensor<f32> loc(#loc5270)
    %4138 = stablehlo.sqrt %4137 : tensor<f32> loc(#loc5249)
    %4139 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4140 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4141 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4142 = call @nan_to_num_1178(%4138, %4139, %4140, %4141) : (tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5257)
    %4143 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4144 = stablehlo.divide %4142, %4143 : tensor<f32> loc(#loc5253)
    %4145 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4146 = stablehlo.maximum %4145, %4144 : tensor<f32> loc(#loc5271)
    %4147 = stablehlo.broadcast_in_dim %4146, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5272)
    %4148 = stablehlo.divide %4084, %4147 : tensor<2048x2048xf32> loc(#loc5272)
    %4149 = stablehlo.multiply %4092, %4092 : tensor<51200xf32> loc(#loc5268)
    %4150 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5269)
    %4151 = stablehlo.reduce(%4149 init: %4150) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5269)
      stablehlo.return %29929 : tensor<f32> loc(#loc5269)
    } loc(#loc5269)
    %4152 = stablehlo.constant dense<5.120000e+04> : tensor<f32> loc(#loc)
    %4153 = stablehlo.divide %4151, %4152 : tensor<f32> loc(#loc5270)
    %4154 = stablehlo.sqrt %4153 : tensor<f32> loc(#loc5249)
    %4155 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4156 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4157 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4158 = call @nan_to_num_1183(%4154, %4155, %4156, %4157) : (tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5257)
    %4159 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4160 = stablehlo.divide %4158, %4159 : tensor<f32> loc(#loc5253)
    %4161 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4162 = stablehlo.maximum %4161, %4160 : tensor<f32> loc(#loc5271)
    %4163 = stablehlo.broadcast_in_dim %4162, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5272)
    %4164 = stablehlo.divide %4092, %4163 : tensor<51200xf32> loc(#loc5272)
    %4165 = stablehlo.multiply %4100, %4100 : tensor<2048x51200xf32> loc(#loc5268)
    %4166 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5274)
    %4167 = stablehlo.reduce(%4165 init: %4166) across dimensions = [1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<2048xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5274)
      stablehlo.return %29929 : tensor<f32> loc(#loc5274)
    } loc(#loc5274)
    %4168 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %4169 = stablehlo.reduce(%4167 init: %4168) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %4170 = stablehlo.constant dense<0x4CC80000> : tensor<f32> loc(#loc)
    %4171 = stablehlo.divide %4169, %4170 : tensor<f32> loc(#loc5253)
    %4172 = stablehlo.sqrt %4171 : tensor<f32> loc(#loc5249)
    %4173 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4174 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4175 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4176 = call @nan_to_num_1188(%4172, %4173, %4174, %4175) : (tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5257)
    %4177 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4178 = stablehlo.divide %4176, %4177 : tensor<f32> loc(#loc5253)
    %4179 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4180 = stablehlo.maximum %4179, %4178 : tensor<f32> loc(#loc5271)
    %4181 = stablehlo.broadcast_in_dim %4180, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5272)
    %4182 = stablehlo.divide %4100, %4181 : tensor<2048x51200xf32> loc(#loc5272)
    %4183 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %4184 = stablehlo.broadcast_in_dim %4183, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5275)
    %4185 = stablehlo.multiply %4184, %arg1 : tensor<2048xf32> loc(#loc5275)
    %4186 = stablehlo.add %4116, %4185 : tensor<2048xf32> loc(#loc5263)
    %4187 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %4188 = stablehlo.broadcast_in_dim %4187, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5275)
    %4189 = stablehlo.multiply %4188, %arg2 : tensor<2048xf32> loc(#loc5275)
    %4190 = stablehlo.add %4132, %4189 : tensor<2048xf32> loc(#loc5263)
    %4191 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %4192 = stablehlo.broadcast_in_dim %4191, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5276)
    %4193 = stablehlo.multiply %4192, %arg3 : tensor<2048x2048xf32> loc(#loc5276)
    %4194 = stablehlo.add %4148, %4193 : tensor<2048x2048xf32> loc(#loc5263)
    %4195 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %4196 = stablehlo.broadcast_in_dim %4195, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5275)
    %4197 = stablehlo.multiply %4196, %arg4 : tensor<51200xf32> loc(#loc5275)
    %4198 = stablehlo.add %4164, %4197 : tensor<51200xf32> loc(#loc5263)
    %4199 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %4200 = stablehlo.broadcast_in_dim %4199, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5276)
    %4201 = stablehlo.multiply %4200, %arg5 : tensor<2048x51200xf32> loc(#loc5276)
    %4202 = stablehlo.add %4182, %4201 : tensor<2048x51200xf32> loc(#loc5263)
    %4203 = stablehlo.convert %arg184 : (tensor<i32>) -> tensor<f32> loc(#loc5225)
    %4204 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4205 = stablehlo.subtract %4203, %4204 : tensor<f32> loc(#loc5226)
    %4206 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4207 = stablehlo.divide %4205, %4206 : tensor<f32> loc(#loc5227)
    %4208 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4209 = stablehlo.multiply %4207, %4208 : tensor<f32> loc(#loc5228)
    %4210 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4211 = stablehlo.add %4210, %4209 : tensor<f32> loc(#loc5229)
    %4212 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4213 = stablehlo.compare  LT, %4203, %4212,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5230)
    %4214 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4215 = stablehlo.compare  GE, %4203, %4214,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5231)
    %4216 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4217 = call @_where_1193(%4215, %4216, %4211) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %4218 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4219 = call @_where_1194(%4213, %4218, %4217) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %4220 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %4221 = stablehlo.compare  LT, %arg184, %4220,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5277)
    %4222 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %4223 = stablehlo.subtract %arg184, %4222 : tensor<i32> loc(#loc5234)
    %4224 = stablehlo.convert %4223 : (tensor<i32>) -> tensor<f32> loc(#loc5225)
    %4225 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4226 = stablehlo.subtract %4224, %4225 : tensor<f32> loc(#loc5226)
    %4227 = stablehlo.constant dense<4.999990e+05> : tensor<f32> loc(#loc)
    %4228 = stablehlo.divide %4226, %4227 : tensor<f32> loc(#loc5227)
    %4229 = stablehlo.constant dense<3.14159274> : tensor<f32> loc(#loc)
    %4230 = stablehlo.multiply %4228, %4229 : tensor<f32> loc(#loc5228)
    %4231 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4232 = stablehlo.add %4231, %4230 : tensor<f32> loc(#loc5229)
    %4233 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %4234 = stablehlo.compare  LT, %4224, %4233,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5230)
    %4235 = stablehlo.constant dense<4.999990e+05> : tensor<f32> loc(#loc)
    %4236 = stablehlo.compare  GE, %4224, %4235,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5231)
    %4237 = stablehlo.constant dense<3.14159274> : tensor<f32> loc(#loc)
    %4238 = call @_where_1195(%4236, %4237, %4232) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %4239 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %4240 = call @_where_1196(%4234, %4239, %4238) : (tensor<i1>, tensor<i32>, tensor<f32>) -> tensor<f32> loc(#loc5232)
    %4241 = stablehlo.cosine %4240 : tensor<f32> loc(#loc5235)
    %4242 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4243 = stablehlo.add %4242, %4241 : tensor<f32> loc(#loc5229)
    %4244 = stablehlo.constant dense<4.500000e-01> : tensor<f32> loc(#loc)
    %4245 = stablehlo.multiply %4244, %4243 : tensor<f32> loc(#loc5236)
    %4246 = stablehlo.constant dense<1.000000e-01> : tensor<f32> loc(#loc)
    %4247 = stablehlo.add %4246, %4245 : tensor<f32> loc(#loc5229)
    %4248 = call @_where_1197(%4221, %4219, %4247) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc5237)
    %4249 = stablehlo.constant dense<6.000000e-04> : tensor<f32> loc(#loc)
    %4250 = stablehlo.multiply %4248, %4249 : tensor<f32> loc(#loc5228)
    %4251 = stablehlo.constant dense<-1.000000e+00> : tensor<f32> loc(#loc)
    %4252 = stablehlo.multiply %4251, %4250 : tensor<f32> loc(#loc5236)
    %4253 = stablehlo.broadcast_in_dim %4252, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %4254 = stablehlo.multiply %4253, %4186 : tensor<2048xf32> loc(#loc5262)
    %4255 = stablehlo.broadcast_in_dim %4252, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5262)
    %4256 = stablehlo.multiply %4255, %4190 : tensor<2048xf32> loc(#loc5262)
    %4257 = stablehlo.broadcast_in_dim %4252, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5262)
    %4258 = stablehlo.multiply %4257, %4194 : tensor<2048x2048xf32> loc(#loc5262)
    %4259 = stablehlo.broadcast_in_dim %4252, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5262)
    %4260 = stablehlo.multiply %4259, %4198 : tensor<51200xf32> loc(#loc5262)
    %4261 = stablehlo.broadcast_in_dim %4252, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5262)
    %4262 = stablehlo.multiply %4261, %4202 : tensor<2048x51200xf32> loc(#loc5262)
    %4263 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %4264 = stablehlo.add %arg184, %4263 : tensor<i32> loc(#loc5256)
    %4265 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %4266 = stablehlo.add %arg195, %4265 : tensor<i32> loc(#loc5256)
    %4267 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %4268 = stablehlo.broadcast_in_dim %4267, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5256)
    %4269 = stablehlo.add %arg196, %4268 : tensor<1xi32> loc(#loc5256)
    %4270 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %4271 = stablehlo.broadcast_in_dim %4270, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5256)
    %4272 = stablehlo.add %arg197, %4271 : tensor<1xi32> loc(#loc5256)
    %4273 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4274 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4275 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4276 = call @nan_to_num_1198(%2711, %4273, %4274, %4275) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4277 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4278 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4279 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4280 = call @nan_to_num_1203(%2713, %4277, %4278, %4279) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4281 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4282 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4283 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4284 = call @nan_to_num_1208(%2715, %4281, %4282, %4283) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4285 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4286 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4287 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4288 = call @nan_to_num_1213(%2717, %4285, %4286, %4287) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4289 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4290 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4291 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4292 = call @nan_to_num_1218(%2719, %4289, %4290, %4291) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4293 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4294 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4295 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4296 = call @nan_to_num_1223(%2721, %4293, %4294, %4295) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4297 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4298 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4299 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4300 = call @nan_to_num_1228(%2723, %4297, %4298, %4299) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4301 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4302 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4303 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4304 = call @nan_to_num_1233(%2725, %4301, %4302, %4303) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4305 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4306 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4307 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4308 = call @nan_to_num_1238(%2727, %4305, %4306, %4307) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4309 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4310 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4311 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4312 = call @nan_to_num_1243(%2729, %4309, %4310, %4311) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4313 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4314 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4315 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4316 = call @nan_to_num_1248(%2731, %4313, %4314, %4315) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4317 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4318 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4319 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4320 = call @nan_to_num_1253(%2733, %4317, %4318, %4319) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4321 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4322 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4323 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4324 = call @nan_to_num_1258(%2735, %4321, %4322, %4323) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4325 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4326 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4327 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4328 = call @nan_to_num_1263(%2737, %4325, %4326, %4327) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4329 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4330 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4331 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4332 = call @nan_to_num_1268(%2739, %4329, %4330, %4331) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4333 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4334 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4335 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4336 = call @nan_to_num_1273(%2741, %4333, %4334, %4335) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4337 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4338 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4339 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4340 = call @nan_to_num_1278(%2743, %4337, %4338, %4339) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4341 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4342 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4343 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4344 = call @nan_to_num_1283(%2745, %4341, %4342, %4343) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4345 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4346 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4347 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4348 = call @nan_to_num_1288(%2747, %4345, %4346, %4347) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4349 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4350 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4351 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4352 = call @nan_to_num_1293(%2749, %4349, %4350, %4351) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4353 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4354 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4355 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4356 = call @nan_to_num_1298(%2751, %4353, %4354, %4355) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4357 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4358 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4359 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4360 = call @nan_to_num_1303(%2753, %4357, %4358, %4359) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4361 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4362 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4363 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4364 = call @nan_to_num_1308(%2755, %4361, %4362, %4363) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4365 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4366 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4367 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4368 = call @nan_to_num_1313(%2757, %4365, %4366, %4367) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4369 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4370 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4371 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4372 = call @nan_to_num_1318(%2759, %4369, %4370, %4371) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4373 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4374 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4375 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4376 = call @nan_to_num_1323(%2761, %4373, %4374, %4375) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4377 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4378 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4379 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4380 = call @nan_to_num_1328(%2763, %4377, %4378, %4379) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4381 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4382 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4383 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4384 = call @nan_to_num_1333(%2765, %4381, %4382, %4383) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4385 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4386 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4387 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4388 = call @nan_to_num_1338(%2767, %4385, %4386, %4387) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4389 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4390 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4391 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4392 = call @nan_to_num_1343(%2769, %4389, %4390, %4391) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4393 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4394 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4395 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4396 = call @nan_to_num_1348(%2771, %4393, %4394, %4395) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4397 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4398 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4399 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4400 = call @nan_to_num_1353(%2773, %4397, %4398, %4399) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4401 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4402 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4403 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4404 = call @nan_to_num_1358(%2775, %4401, %4402, %4403) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4405 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4406 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4407 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4408 = call @nan_to_num_1363(%2777, %4405, %4406, %4407) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4409 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4410 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4411 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4412 = call @nan_to_num_1368(%2779, %4409, %4410, %4411) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4413 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4414 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4415 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4416 = call @nan_to_num_1373(%2781, %4413, %4414, %4415) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4417 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4418 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4419 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4420 = call @nan_to_num_1378(%2783, %4417, %4418, %4419) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4421 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4422 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4423 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4424 = call @nan_to_num_1383(%2785, %4421, %4422, %4423) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4425 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4426 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4427 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4428 = call @nan_to_num_1388(%2787, %4425, %4426, %4427) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4429 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4430 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4431 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4432 = call @nan_to_num_1393(%2789, %4429, %4430, %4431) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4433 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4434 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4435 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4436 = call @nan_to_num_1398(%2791, %4433, %4434, %4435) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4437 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4438 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4439 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4440 = call @nan_to_num_1403(%2793, %4437, %4438, %4439) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4441 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4442 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4443 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4444 = call @nan_to_num_1408(%2795, %4441, %4442, %4443) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4445 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4446 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4447 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4448 = call @nan_to_num_1413(%2797, %4445, %4446, %4447) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4449 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4450 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4451 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4452 = call @nan_to_num_1418(%2799, %4449, %4450, %4451) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4453 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4454 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4455 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4456 = call @nan_to_num_1423(%2801, %4453, %4454, %4455) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4457 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4458 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4459 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4460 = call @nan_to_num_1428(%2803, %4457, %4458, %4459) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4461 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4462 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4463 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4464 = call @nan_to_num_1433(%2805, %4461, %4462, %4463) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4465 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4466 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4467 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4468 = call @nan_to_num_1438(%2807, %4465, %4466, %4467) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4469 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4470 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4471 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4472 = call @nan_to_num_1443(%2809, %4469, %4470, %4471) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4473 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4474 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4475 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4476 = call @nan_to_num_1448(%2811, %4473, %4474, %4475) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4477 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4478 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4479 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4480 = call @nan_to_num_1453(%2813, %4477, %4478, %4479) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4481 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4482 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4483 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4484 = call @nan_to_num_1458(%2815, %4481, %4482, %4483) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4485 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4486 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4487 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4488 = call @nan_to_num_1463(%2817, %4485, %4486, %4487) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4489 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4490 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4491 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4492 = call @nan_to_num_1468(%2819, %4489, %4490, %4491) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4493 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4494 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4495 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4496 = call @nan_to_num_1473(%2821, %4493, %4494, %4495) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4497 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4498 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4499 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4500 = call @nan_to_num_1478(%2823, %4497, %4498, %4499) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4501 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4502 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4503 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4504 = call @nan_to_num_1483(%2825, %4501, %4502, %4503) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4505 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4506 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4507 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4508 = call @nan_to_num_1488(%2827, %4505, %4506, %4507) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4509 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4510 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4511 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4512 = call @nan_to_num_1493(%2829, %4509, %4510, %4511) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4513 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4514 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4515 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4516 = call @nan_to_num_1498(%2831, %4513, %4514, %4515) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4517 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4518 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4519 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4520 = call @nan_to_num_1503(%2833, %4517, %4518, %4519) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4521 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4522 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4523 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4524 = call @nan_to_num_1508(%2835, %4521, %4522, %4523) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4525 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4526 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4527 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4528 = call @nan_to_num_1513(%2837, %4525, %4526, %4527) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4529 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4530 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4531 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4532 = call @nan_to_num_1518(%2839, %4529, %4530, %4531) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4533 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4534 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4535 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4536 = call @nan_to_num_1523(%2841, %4533, %4534, %4535) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4537 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4538 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4539 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4540 = call @nan_to_num_1528(%2843, %4537, %4538, %4539) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4541 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4542 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4543 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4544 = call @nan_to_num_1533(%2845, %4541, %4542, %4543) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4545 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4546 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4547 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4548 = call @nan_to_num_1538(%2847, %4545, %4546, %4547) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4549 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4550 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4551 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4552 = call @nan_to_num_1543(%2849, %4549, %4550, %4551) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4553 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4554 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4555 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4556 = call @nan_to_num_1548(%2851, %4553, %4554, %4555) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4557 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4558 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4559 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4560 = call @nan_to_num_1553(%2853, %4557, %4558, %4559) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4561 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4562 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4563 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4564 = call @nan_to_num_1558(%2855, %4561, %4562, %4563) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4565 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4566 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4567 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4568 = call @nan_to_num_1563(%2857, %4565, %4566, %4567) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4569 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4570 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4571 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4572 = call @nan_to_num_1568(%2859, %4569, %4570, %4571) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4573 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4574 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4575 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4576 = call @nan_to_num_1573(%2861, %4573, %4574, %4575) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4577 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4578 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4579 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4580 = call @nan_to_num_1578(%2863, %4577, %4578, %4579) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4581 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4582 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4583 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4584 = call @nan_to_num_1583(%2865, %4581, %4582, %4583) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4585 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4586 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4587 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4588 = call @nan_to_num_1588(%2867, %4585, %4586, %4587) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4589 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4590 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4591 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4592 = call @nan_to_num_1593(%2869, %4589, %4590, %4591) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4593 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4594 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4595 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4596 = call @nan_to_num_1598(%2871, %4593, %4594, %4595) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4597 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4598 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4599 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4600 = call @nan_to_num_1603(%2873, %4597, %4598, %4599) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4601 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4602 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4603 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4604 = call @nan_to_num_1608(%2875, %4601, %4602, %4603) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4605 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4606 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4607 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4608 = call @nan_to_num_1613(%2877, %4605, %4606, %4607) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4609 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4610 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4611 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4612 = call @nan_to_num_1618(%2879, %4609, %4610, %4611) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4613 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4614 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4615 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4616 = call @nan_to_num_1623(%2881, %4613, %4614, %4615) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4617 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4618 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4619 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4620 = call @nan_to_num_1628(%2883, %4617, %4618, %4619) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4621 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4622 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4623 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4624 = call @nan_to_num_1633(%2885, %4621, %4622, %4623) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4625 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4626 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4627 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4628 = call @nan_to_num_1638(%2887, %4625, %4626, %4627) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4629 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4630 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4631 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4632 = call @nan_to_num_1643(%2889, %4629, %4630, %4631) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4633 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4634 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4635 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4636 = call @nan_to_num_1648(%2891, %4633, %4634, %4635) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4637 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4638 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4639 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4640 = call @nan_to_num_1653(%2893, %4637, %4638, %4639) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4641 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4642 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4643 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4644 = call @nan_to_num_1658(%2895, %4641, %4642, %4643) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4645 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4646 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4647 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4648 = call @nan_to_num_1663(%2897, %4645, %4646, %4647) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4649 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4650 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4651 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4652 = call @nan_to_num_1668(%2899, %4649, %4650, %4651) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4653 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4654 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4655 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4656 = call @nan_to_num_1673(%2901, %4653, %4654, %4655) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4657 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4658 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4659 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4660 = call @nan_to_num_1678(%2903, %4657, %4658, %4659) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4661 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4662 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4663 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4664 = call @nan_to_num_1683(%2905, %4661, %4662, %4663) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4665 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4666 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4667 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4668 = call @nan_to_num_1688(%2907, %4665, %4666, %4667) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4669 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4670 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4671 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4672 = call @nan_to_num_1693(%2909, %4669, %4670, %4671) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4673 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4674 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4675 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4676 = call @nan_to_num_1698(%2911, %4673, %4674, %4675) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4677 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4678 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4679 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4680 = call @nan_to_num_1703(%2913, %4677, %4678, %4679) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4681 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4682 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4683 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4684 = call @nan_to_num_1708(%2915, %4681, %4682, %4683) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4685 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4686 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4687 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4688 = call @nan_to_num_1713(%2917, %4685, %4686, %4687) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4689 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4690 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4691 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4692 = call @nan_to_num_1718(%2919, %4689, %4690, %4691) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4693 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4694 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4695 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4696 = call @nan_to_num_1723(%2921, %4693, %4694, %4695) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4697 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4698 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4699 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4700 = call @nan_to_num_1728(%2923, %4697, %4698, %4699) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4701 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4702 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4703 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4704 = call @nan_to_num_1733(%2925, %4701, %4702, %4703) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4705 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4706 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4707 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4708 = call @nan_to_num_1738(%2927, %4705, %4706, %4707) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4709 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4710 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4711 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4712 = call @nan_to_num_1743(%2929, %4709, %4710, %4711) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4713 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4714 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4715 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4716 = call @nan_to_num_1748(%2931, %4713, %4714, %4715) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4717 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4718 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4719 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4720 = call @nan_to_num_1753(%2933, %4717, %4718, %4719) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4721 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4722 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4723 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4724 = call @nan_to_num_1758(%2935, %4721, %4722, %4723) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4725 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4726 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4727 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4728 = call @nan_to_num_1763(%2937, %4725, %4726, %4727) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4729 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4730 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4731 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4732 = call @nan_to_num_1768(%2939, %4729, %4730, %4731) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4733 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4734 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4735 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4736 = call @nan_to_num_1773(%2941, %4733, %4734, %4735) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4737 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4738 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4739 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4740 = call @nan_to_num_1778(%2943, %4737, %4738, %4739) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4741 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4742 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4743 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4744 = call @nan_to_num_1783(%2945, %4741, %4742, %4743) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4745 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4746 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4747 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4748 = call @nan_to_num_1788(%2947, %4745, %4746, %4747) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4749 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4750 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4751 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4752 = call @nan_to_num_1793(%2949, %4749, %4750, %4751) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4753 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4754 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4755 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4756 = call @nan_to_num_1798(%2951, %4753, %4754, %4755) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4757 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4758 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4759 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4760 = call @nan_to_num_1803(%2953, %4757, %4758, %4759) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4761 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4762 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4763 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4764 = call @nan_to_num_1808(%2955, %4761, %4762, %4763) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4765 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4766 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4767 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4768 = call @nan_to_num_1813(%2957, %4765, %4766, %4767) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4769 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4770 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4771 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4772 = call @nan_to_num_1818(%2959, %4769, %4770, %4771) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4773 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4774 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4775 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4776 = call @nan_to_num_1823(%2961, %4773, %4774, %4775) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4777 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4778 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4779 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4780 = call @nan_to_num_1828(%2963, %4777, %4778, %4779) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4781 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4782 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4783 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4784 = call @nan_to_num_1833(%2965, %4781, %4782, %4783) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4785 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4786 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4787 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4788 = call @nan_to_num_1838(%2967, %4785, %4786, %4787) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4789 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4790 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4791 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4792 = call @nan_to_num_1843(%2969, %4789, %4790, %4791) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4793 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4794 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4795 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4796 = call @nan_to_num_1848(%2971, %4793, %4794, %4795) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4797 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4798 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4799 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4800 = call @nan_to_num_1853(%2973, %4797, %4798, %4799) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4801 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4802 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4803 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4804 = call @nan_to_num_1858(%2975, %4801, %4802, %4803) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4805 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4806 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4807 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4808 = call @nan_to_num_1863(%2977, %4805, %4806, %4807) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4809 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4810 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4811 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4812 = call @nan_to_num_1868(%2979, %4809, %4810, %4811) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4813 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4814 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4815 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4816 = call @nan_to_num_1873(%2981, %4813, %4814, %4815) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4817 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4818 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4819 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4820 = call @nan_to_num_1878(%2983, %4817, %4818, %4819) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4821 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4822 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4823 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4824 = call @nan_to_num_1883(%2985, %4821, %4822, %4823) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4825 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4826 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4827 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4828 = call @nan_to_num_1888(%2987, %4825, %4826, %4827) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4829 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4830 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4831 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4832 = call @nan_to_num_1893(%2989, %4829, %4830, %4831) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4833 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4834 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4835 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4836 = call @nan_to_num_1898(%2991, %4833, %4834, %4835) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4837 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4838 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4839 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4840 = call @nan_to_num_1903(%2993, %4837, %4838, %4839) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4841 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4842 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4843 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4844 = call @nan_to_num_1908(%2995, %4841, %4842, %4843) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4845 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4846 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4847 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4848 = call @nan_to_num_1913(%2997, %4845, %4846, %4847) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4849 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4850 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4851 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4852 = call @nan_to_num_1918(%2999, %4849, %4850, %4851) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4853 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4854 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4855 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4856 = call @nan_to_num_1923(%3001, %4853, %4854, %4855) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4857 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4858 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4859 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4860 = call @nan_to_num_1928(%3003, %4857, %4858, %4859) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4861 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4862 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4863 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4864 = call @nan_to_num_1933(%3005, %4861, %4862, %4863) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4865 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4866 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4867 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4868 = call @nan_to_num_1938(%3007, %4865, %4866, %4867) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4869 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4870 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4871 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4872 = call @nan_to_num_1943(%3009, %4869, %4870, %4871) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4873 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4874 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4875 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4876 = call @nan_to_num_1948(%3011, %4873, %4874, %4875) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4877 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4878 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4879 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4880 = call @nan_to_num_1953(%3013, %4877, %4878, %4879) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4881 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4882 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4883 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4884 = call @nan_to_num_1958(%3015, %4881, %4882, %4883) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4885 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4886 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4887 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4888 = call @nan_to_num_1963(%3017, %4885, %4886, %4887) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4889 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4890 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4891 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4892 = call @nan_to_num_1968(%3019, %4889, %4890, %4891) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4893 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4894 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4895 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4896 = call @nan_to_num_1973(%3021, %4893, %4894, %4895) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4897 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4898 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4899 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4900 = call @nan_to_num_1978(%3023, %4897, %4898, %4899) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4901 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4902 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4903 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4904 = call @nan_to_num_1983(%3025, %4901, %4902, %4903) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4905 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4906 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4907 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4908 = call @nan_to_num_1988(%3027, %4905, %4906, %4907) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4909 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4910 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4911 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4912 = call @nan_to_num_1993(%3029, %4909, %4910, %4911) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4913 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4914 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4915 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4916 = call @nan_to_num_1998(%3031, %4913, %4914, %4915) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4917 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4918 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4919 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4920 = call @nan_to_num_2003(%3033, %4917, %4918, %4919) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4921 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4922 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4923 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4924 = call @nan_to_num_2008(%3035, %4921, %4922, %4923) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4925 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4926 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4927 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4928 = call @nan_to_num_2013(%3037, %4925, %4926, %4927) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4929 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4930 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4931 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4932 = call @nan_to_num_2018(%3039, %4929, %4930, %4931) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4933 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4934 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4935 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4936 = call @nan_to_num_2023(%3041, %4933, %4934, %4935) : (tensor<1x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192xf32> loc(#loc5257)
    %4937 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4938 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4939 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4940 = call @nan_to_num_2028(%3043, %4937, %4938, %4939) : (tensor<1x2048x8192xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5257)
    %4941 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4942 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4943 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4944 = call @nan_to_num_2033(%3045, %4941, %4942, %4943) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4945 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4946 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4947 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4948 = call @nan_to_num_2038(%3047, %4945, %4946, %4947) : (tensor<1x8192x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5257)
    %4949 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4950 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4951 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4952 = call @nan_to_num_2043(%3049, %4949, %4950, %4951) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4953 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4954 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4955 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4956 = call @nan_to_num_2048(%3051, %4953, %4954, %4955) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4957 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4958 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4959 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4960 = call @nan_to_num_2053(%3053, %4957, %4958, %4959) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4961 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4962 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4963 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4964 = call @nan_to_num_2058(%3055, %4961, %4962, %4963) : (tensor<1x2048xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048xf32> loc(#loc5257)
    %4965 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4966 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4967 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4968 = call @nan_to_num_2063(%3057, %4965, %4966, %4967) : (tensor<1x3x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5257)
    %4969 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4970 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4971 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4972 = call @nan_to_num_2068(%3059, %4969, %4970, %4971) : (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x64xf32> loc(#loc5257)
    %4973 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4974 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4975 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %4976 = call @nan_to_num_2073(%3061, %4973, %4974, %4975) : (tensor<1x2048x32x64xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5257)
    %4977 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %4978 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4979 = stablehlo.broadcast_in_dim %4978, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %4980 = stablehlo.add %4977, %4979 : tensor<1xf32> loc(#loc5259)
    %4981 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4982 = stablehlo.broadcast_in_dim %4981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %4983 = stablehlo.subtract %4980, %4982 : tensor<1xf32> loc(#loc5260)
    %4984 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %4985 = stablehlo.broadcast_in_dim %4984, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %4986 = stablehlo.power %4985, %4983 : tensor<1xf32> loc(#loc5261)
    %4987 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4988 = stablehlo.broadcast_in_dim %4987, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %4989 = stablehlo.subtract %4988, %4986 : tensor<1xf32> loc(#loc5260)
    %4990 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %4991 = stablehlo.broadcast_in_dim %4990, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %4992 = stablehlo.multiply %4991, %4989 : tensor<1xf32> loc(#loc5236)
    %4993 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %4994 = stablehlo.broadcast_in_dim %4993, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %4995 = stablehlo.power %4994, %4980 : tensor<1xf32> loc(#loc5261)
    %4996 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %4997 = stablehlo.broadcast_in_dim %4996, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %4998 = stablehlo.subtract %4997, %4995 : tensor<1xf32> loc(#loc5260)
    %4999 = stablehlo.divide %4992, %4998 : tensor<1xf32> loc(#loc5253)
    %5000 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5001 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5002 = stablehlo.broadcast_in_dim %5001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5003 = stablehlo.add %5000, %5002 : tensor<1xf32> loc(#loc5259)
    %5004 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5005 = stablehlo.broadcast_in_dim %5004, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5006 = stablehlo.subtract %5003, %5005 : tensor<1xf32> loc(#loc5260)
    %5007 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5008 = stablehlo.broadcast_in_dim %5007, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5009 = stablehlo.power %5008, %5006 : tensor<1xf32> loc(#loc5261)
    %5010 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5011 = stablehlo.broadcast_in_dim %5010, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5012 = stablehlo.subtract %5011, %5009 : tensor<1xf32> loc(#loc5260)
    %5013 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5014 = stablehlo.broadcast_in_dim %5013, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5015 = stablehlo.multiply %5014, %5012 : tensor<1xf32> loc(#loc5236)
    %5016 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5017 = stablehlo.broadcast_in_dim %5016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5018 = stablehlo.power %5017, %5003 : tensor<1xf32> loc(#loc5261)
    %5019 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5020 = stablehlo.broadcast_in_dim %5019, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5021 = stablehlo.subtract %5020, %5018 : tensor<1xf32> loc(#loc5260)
    %5022 = stablehlo.divide %5015, %5021 : tensor<1xf32> loc(#loc5253)
    %5023 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5024 = stablehlo.broadcast_in_dim %5023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5025 = stablehlo.subtract %5024, %4999 : tensor<1xf32> loc(#loc5260)
    %5026 = stablehlo.broadcast_in_dim %5025, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5027 = stablehlo.broadcast_in_dim %5026, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5028 = stablehlo.multiply %5027, %4276 : tensor<1x8192xf32> loc(#loc5262)
    %5029 = stablehlo.broadcast_in_dim %4999, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5030 = stablehlo.broadcast_in_dim %5029, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5031 = stablehlo.multiply %5030, %arg199 : tensor<1x8192xf32> loc(#loc5262)
    %5032 = stablehlo.add %5028, %5031 : tensor<1x8192xf32> loc(#loc5263)
    %5033 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5034 = stablehlo.broadcast_in_dim %5033, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5035 = stablehlo.subtract %5034, %5022 : tensor<1xf32> loc(#loc5260)
    %5036 = stablehlo.multiply %4276, %4276 : tensor<1x8192xf32> loc(#loc5264)
    %5037 = stablehlo.broadcast_in_dim %5035, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5038 = stablehlo.broadcast_in_dim %5037, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5039 = stablehlo.multiply %5038, %5036 : tensor<1x8192xf32> loc(#loc5262)
    %5040 = stablehlo.broadcast_in_dim %5022, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5041 = stablehlo.broadcast_in_dim %5040, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5042 = stablehlo.multiply %5041, %arg375 : tensor<1x8192xf32> loc(#loc5262)
    %5043 = stablehlo.add %5039, %5042 : tensor<1x8192xf32> loc(#loc5263)
    %5044 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5045 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5046 = stablehlo.broadcast_in_dim %5045, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5047 = stablehlo.add %5044, %5046 : tensor<1xf32> loc(#loc5259)
    %5048 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5049 = stablehlo.broadcast_in_dim %5048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5050 = stablehlo.subtract %5047, %5049 : tensor<1xf32> loc(#loc5260)
    %5051 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5052 = stablehlo.broadcast_in_dim %5051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5053 = stablehlo.power %5052, %5050 : tensor<1xf32> loc(#loc5261)
    %5054 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5055 = stablehlo.broadcast_in_dim %5054, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5056 = stablehlo.subtract %5055, %5053 : tensor<1xf32> loc(#loc5260)
    %5057 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5058 = stablehlo.broadcast_in_dim %5057, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5059 = stablehlo.multiply %5058, %5056 : tensor<1xf32> loc(#loc5236)
    %5060 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5061 = stablehlo.broadcast_in_dim %5060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5062 = stablehlo.power %5061, %5047 : tensor<1xf32> loc(#loc5261)
    %5063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5064 = stablehlo.broadcast_in_dim %5063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5065 = stablehlo.subtract %5064, %5062 : tensor<1xf32> loc(#loc5260)
    %5066 = stablehlo.divide %5059, %5065 : tensor<1xf32> loc(#loc5253)
    %5067 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5069 = stablehlo.broadcast_in_dim %5068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5070 = stablehlo.add %5067, %5069 : tensor<1xf32> loc(#loc5259)
    %5071 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5072 = stablehlo.broadcast_in_dim %5071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5073 = stablehlo.subtract %5070, %5072 : tensor<1xf32> loc(#loc5260)
    %5074 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5075 = stablehlo.broadcast_in_dim %5074, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5076 = stablehlo.power %5075, %5073 : tensor<1xf32> loc(#loc5261)
    %5077 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5078 = stablehlo.broadcast_in_dim %5077, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5079 = stablehlo.subtract %5078, %5076 : tensor<1xf32> loc(#loc5260)
    %5080 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5081 = stablehlo.broadcast_in_dim %5080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5082 = stablehlo.multiply %5081, %5079 : tensor<1xf32> loc(#loc5236)
    %5083 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5084 = stablehlo.broadcast_in_dim %5083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5085 = stablehlo.power %5084, %5070 : tensor<1xf32> loc(#loc5261)
    %5086 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5087 = stablehlo.broadcast_in_dim %5086, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5088 = stablehlo.subtract %5087, %5085 : tensor<1xf32> loc(#loc5260)
    %5089 = stablehlo.divide %5082, %5088 : tensor<1xf32> loc(#loc5253)
    %5090 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5091 = stablehlo.broadcast_in_dim %5090, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5092 = stablehlo.subtract %5091, %5066 : tensor<1xf32> loc(#loc5260)
    %5093 = stablehlo.broadcast_in_dim %5092, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5094 = stablehlo.broadcast_in_dim %5093, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5095 = stablehlo.multiply %5094, %4280 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5096 = stablehlo.broadcast_in_dim %5066, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5097 = stablehlo.broadcast_in_dim %5096, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5098 = stablehlo.multiply %5097, %arg200 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5099 = stablehlo.add %5095, %5098 : tensor<1x2048x8192xf32> loc(#loc5263)
    %5100 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5101 = stablehlo.broadcast_in_dim %5100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5102 = stablehlo.subtract %5101, %5089 : tensor<1xf32> loc(#loc5260)
    %5103 = stablehlo.multiply %4280, %4280 : tensor<1x2048x8192xf32> loc(#loc5264)
    %5104 = stablehlo.broadcast_in_dim %5102, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5105 = stablehlo.broadcast_in_dim %5104, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5106 = stablehlo.multiply %5105, %5103 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5107 = stablehlo.broadcast_in_dim %5089, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5108 = stablehlo.broadcast_in_dim %5107, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5109 = stablehlo.multiply %5108, %arg376 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5110 = stablehlo.add %5106, %5109 : tensor<1x2048x8192xf32> loc(#loc5263)
    %5111 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5112 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5113 = stablehlo.broadcast_in_dim %5112, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5114 = stablehlo.add %5111, %5113 : tensor<1xf32> loc(#loc5259)
    %5115 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5116 = stablehlo.broadcast_in_dim %5115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5117 = stablehlo.subtract %5114, %5116 : tensor<1xf32> loc(#loc5260)
    %5118 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5119 = stablehlo.broadcast_in_dim %5118, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5120 = stablehlo.power %5119, %5117 : tensor<1xf32> loc(#loc5261)
    %5121 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5122 = stablehlo.broadcast_in_dim %5121, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5123 = stablehlo.subtract %5122, %5120 : tensor<1xf32> loc(#loc5260)
    %5124 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5125 = stablehlo.broadcast_in_dim %5124, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5126 = stablehlo.multiply %5125, %5123 : tensor<1xf32> loc(#loc5236)
    %5127 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5128 = stablehlo.broadcast_in_dim %5127, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5129 = stablehlo.power %5128, %5114 : tensor<1xf32> loc(#loc5261)
    %5130 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5131 = stablehlo.broadcast_in_dim %5130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5132 = stablehlo.subtract %5131, %5129 : tensor<1xf32> loc(#loc5260)
    %5133 = stablehlo.divide %5126, %5132 : tensor<1xf32> loc(#loc5253)
    %5134 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5135 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5136 = stablehlo.broadcast_in_dim %5135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5137 = stablehlo.add %5134, %5136 : tensor<1xf32> loc(#loc5259)
    %5138 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5139 = stablehlo.broadcast_in_dim %5138, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5140 = stablehlo.subtract %5137, %5139 : tensor<1xf32> loc(#loc5260)
    %5141 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5142 = stablehlo.broadcast_in_dim %5141, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5143 = stablehlo.power %5142, %5140 : tensor<1xf32> loc(#loc5261)
    %5144 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5145 = stablehlo.broadcast_in_dim %5144, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5146 = stablehlo.subtract %5145, %5143 : tensor<1xf32> loc(#loc5260)
    %5147 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5148 = stablehlo.broadcast_in_dim %5147, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5149 = stablehlo.multiply %5148, %5146 : tensor<1xf32> loc(#loc5236)
    %5150 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5151 = stablehlo.broadcast_in_dim %5150, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5152 = stablehlo.power %5151, %5137 : tensor<1xf32> loc(#loc5261)
    %5153 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5154 = stablehlo.broadcast_in_dim %5153, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5155 = stablehlo.subtract %5154, %5152 : tensor<1xf32> loc(#loc5260)
    %5156 = stablehlo.divide %5149, %5155 : tensor<1xf32> loc(#loc5253)
    %5157 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5158 = stablehlo.broadcast_in_dim %5157, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5159 = stablehlo.subtract %5158, %5133 : tensor<1xf32> loc(#loc5260)
    %5160 = stablehlo.broadcast_in_dim %5159, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5161 = stablehlo.broadcast_in_dim %5160, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5162 = stablehlo.multiply %5161, %4284 : tensor<1x2048xf32> loc(#loc5262)
    %5163 = stablehlo.broadcast_in_dim %5133, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5164 = stablehlo.broadcast_in_dim %5163, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5165 = stablehlo.multiply %5164, %arg201 : tensor<1x2048xf32> loc(#loc5262)
    %5166 = stablehlo.add %5162, %5165 : tensor<1x2048xf32> loc(#loc5263)
    %5167 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5168 = stablehlo.broadcast_in_dim %5167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5169 = stablehlo.subtract %5168, %5156 : tensor<1xf32> loc(#loc5260)
    %5170 = stablehlo.multiply %4284, %4284 : tensor<1x2048xf32> loc(#loc5264)
    %5171 = stablehlo.broadcast_in_dim %5169, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5172 = stablehlo.broadcast_in_dim %5171, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5173 = stablehlo.multiply %5172, %5170 : tensor<1x2048xf32> loc(#loc5262)
    %5174 = stablehlo.broadcast_in_dim %5156, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5175 = stablehlo.broadcast_in_dim %5174, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5176 = stablehlo.multiply %5175, %arg377 : tensor<1x2048xf32> loc(#loc5262)
    %5177 = stablehlo.add %5173, %5176 : tensor<1x2048xf32> loc(#loc5263)
    %5178 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5179 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5180 = stablehlo.broadcast_in_dim %5179, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5181 = stablehlo.add %5178, %5180 : tensor<1xf32> loc(#loc5259)
    %5182 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5183 = stablehlo.broadcast_in_dim %5182, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5184 = stablehlo.subtract %5181, %5183 : tensor<1xf32> loc(#loc5260)
    %5185 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5186 = stablehlo.broadcast_in_dim %5185, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5187 = stablehlo.power %5186, %5184 : tensor<1xf32> loc(#loc5261)
    %5188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5189 = stablehlo.broadcast_in_dim %5188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5190 = stablehlo.subtract %5189, %5187 : tensor<1xf32> loc(#loc5260)
    %5191 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5192 = stablehlo.broadcast_in_dim %5191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5193 = stablehlo.multiply %5192, %5190 : tensor<1xf32> loc(#loc5236)
    %5194 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5195 = stablehlo.broadcast_in_dim %5194, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5196 = stablehlo.power %5195, %5181 : tensor<1xf32> loc(#loc5261)
    %5197 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5198 = stablehlo.broadcast_in_dim %5197, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5199 = stablehlo.subtract %5198, %5196 : tensor<1xf32> loc(#loc5260)
    %5200 = stablehlo.divide %5193, %5199 : tensor<1xf32> loc(#loc5253)
    %5201 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5202 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5203 = stablehlo.broadcast_in_dim %5202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5204 = stablehlo.add %5201, %5203 : tensor<1xf32> loc(#loc5259)
    %5205 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5206 = stablehlo.broadcast_in_dim %5205, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5207 = stablehlo.subtract %5204, %5206 : tensor<1xf32> loc(#loc5260)
    %5208 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5209 = stablehlo.broadcast_in_dim %5208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5210 = stablehlo.power %5209, %5207 : tensor<1xf32> loc(#loc5261)
    %5211 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5212 = stablehlo.broadcast_in_dim %5211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5213 = stablehlo.subtract %5212, %5210 : tensor<1xf32> loc(#loc5260)
    %5214 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5215 = stablehlo.broadcast_in_dim %5214, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5216 = stablehlo.multiply %5215, %5213 : tensor<1xf32> loc(#loc5236)
    %5217 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5218 = stablehlo.broadcast_in_dim %5217, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5219 = stablehlo.power %5218, %5204 : tensor<1xf32> loc(#loc5261)
    %5220 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5221 = stablehlo.broadcast_in_dim %5220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5222 = stablehlo.subtract %5221, %5219 : tensor<1xf32> loc(#loc5260)
    %5223 = stablehlo.divide %5216, %5222 : tensor<1xf32> loc(#loc5253)
    %5224 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5225 = stablehlo.broadcast_in_dim %5224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5226 = stablehlo.subtract %5225, %5200 : tensor<1xf32> loc(#loc5260)
    %5227 = stablehlo.broadcast_in_dim %5226, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5228 = stablehlo.broadcast_in_dim %5227, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5229 = stablehlo.multiply %5228, %4288 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5230 = stablehlo.broadcast_in_dim %5200, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5231 = stablehlo.broadcast_in_dim %5230, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5232 = stablehlo.multiply %5231, %arg202 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5233 = stablehlo.add %5229, %5232 : tensor<1x8192x2048xf32> loc(#loc5263)
    %5234 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5235 = stablehlo.broadcast_in_dim %5234, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5236 = stablehlo.subtract %5235, %5223 : tensor<1xf32> loc(#loc5260)
    %5237 = stablehlo.multiply %4288, %4288 : tensor<1x8192x2048xf32> loc(#loc5264)
    %5238 = stablehlo.broadcast_in_dim %5236, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5239 = stablehlo.broadcast_in_dim %5238, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5240 = stablehlo.multiply %5239, %5237 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5241 = stablehlo.broadcast_in_dim %5223, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5242 = stablehlo.broadcast_in_dim %5241, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5243 = stablehlo.multiply %5242, %arg378 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5244 = stablehlo.add %5240, %5243 : tensor<1x8192x2048xf32> loc(#loc5263)
    %5245 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5246 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5247 = stablehlo.broadcast_in_dim %5246, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5248 = stablehlo.add %5245, %5247 : tensor<1xf32> loc(#loc5259)
    %5249 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5250 = stablehlo.broadcast_in_dim %5249, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5251 = stablehlo.subtract %5248, %5250 : tensor<1xf32> loc(#loc5260)
    %5252 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5253 = stablehlo.broadcast_in_dim %5252, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5254 = stablehlo.power %5253, %5251 : tensor<1xf32> loc(#loc5261)
    %5255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5256 = stablehlo.broadcast_in_dim %5255, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5257 = stablehlo.subtract %5256, %5254 : tensor<1xf32> loc(#loc5260)
    %5258 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5259 = stablehlo.broadcast_in_dim %5258, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5260 = stablehlo.multiply %5259, %5257 : tensor<1xf32> loc(#loc5236)
    %5261 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5262 = stablehlo.broadcast_in_dim %5261, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5263 = stablehlo.power %5262, %5248 : tensor<1xf32> loc(#loc5261)
    %5264 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5265 = stablehlo.broadcast_in_dim %5264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5266 = stablehlo.subtract %5265, %5263 : tensor<1xf32> loc(#loc5260)
    %5267 = stablehlo.divide %5260, %5266 : tensor<1xf32> loc(#loc5253)
    %5268 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5270 = stablehlo.broadcast_in_dim %5269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5271 = stablehlo.add %5268, %5270 : tensor<1xf32> loc(#loc5259)
    %5272 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5273 = stablehlo.broadcast_in_dim %5272, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5274 = stablehlo.subtract %5271, %5273 : tensor<1xf32> loc(#loc5260)
    %5275 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5276 = stablehlo.broadcast_in_dim %5275, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5277 = stablehlo.power %5276, %5274 : tensor<1xf32> loc(#loc5261)
    %5278 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5279 = stablehlo.broadcast_in_dim %5278, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5280 = stablehlo.subtract %5279, %5277 : tensor<1xf32> loc(#loc5260)
    %5281 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5282 = stablehlo.broadcast_in_dim %5281, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5283 = stablehlo.multiply %5282, %5280 : tensor<1xf32> loc(#loc5236)
    %5284 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5285 = stablehlo.broadcast_in_dim %5284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5286 = stablehlo.power %5285, %5271 : tensor<1xf32> loc(#loc5261)
    %5287 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5288 = stablehlo.broadcast_in_dim %5287, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5289 = stablehlo.subtract %5288, %5286 : tensor<1xf32> loc(#loc5260)
    %5290 = stablehlo.divide %5283, %5289 : tensor<1xf32> loc(#loc5253)
    %5291 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5292 = stablehlo.broadcast_in_dim %5291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5293 = stablehlo.subtract %5292, %5267 : tensor<1xf32> loc(#loc5260)
    %5294 = stablehlo.broadcast_in_dim %5293, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5295 = stablehlo.broadcast_in_dim %5294, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5296 = stablehlo.multiply %5295, %4292 : tensor<1x2048xf32> loc(#loc5262)
    %5297 = stablehlo.broadcast_in_dim %5267, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5298 = stablehlo.broadcast_in_dim %5297, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5299 = stablehlo.multiply %5298, %arg203 : tensor<1x2048xf32> loc(#loc5262)
    %5300 = stablehlo.add %5296, %5299 : tensor<1x2048xf32> loc(#loc5263)
    %5301 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5302 = stablehlo.broadcast_in_dim %5301, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5303 = stablehlo.subtract %5302, %5290 : tensor<1xf32> loc(#loc5260)
    %5304 = stablehlo.multiply %4292, %4292 : tensor<1x2048xf32> loc(#loc5264)
    %5305 = stablehlo.broadcast_in_dim %5303, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5306 = stablehlo.broadcast_in_dim %5305, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5307 = stablehlo.multiply %5306, %5304 : tensor<1x2048xf32> loc(#loc5262)
    %5308 = stablehlo.broadcast_in_dim %5290, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5309 = stablehlo.broadcast_in_dim %5308, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5310 = stablehlo.multiply %5309, %arg379 : tensor<1x2048xf32> loc(#loc5262)
    %5311 = stablehlo.add %5307, %5310 : tensor<1x2048xf32> loc(#loc5263)
    %5312 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5313 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5314 = stablehlo.broadcast_in_dim %5313, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5315 = stablehlo.add %5312, %5314 : tensor<1xf32> loc(#loc5259)
    %5316 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5317 = stablehlo.broadcast_in_dim %5316, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5318 = stablehlo.subtract %5315, %5317 : tensor<1xf32> loc(#loc5260)
    %5319 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5320 = stablehlo.broadcast_in_dim %5319, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5321 = stablehlo.power %5320, %5318 : tensor<1xf32> loc(#loc5261)
    %5322 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5323 = stablehlo.broadcast_in_dim %5322, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5324 = stablehlo.subtract %5323, %5321 : tensor<1xf32> loc(#loc5260)
    %5325 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5326 = stablehlo.broadcast_in_dim %5325, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5327 = stablehlo.multiply %5326, %5324 : tensor<1xf32> loc(#loc5236)
    %5328 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5329 = stablehlo.broadcast_in_dim %5328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5330 = stablehlo.power %5329, %5315 : tensor<1xf32> loc(#loc5261)
    %5331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5332 = stablehlo.broadcast_in_dim %5331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5333 = stablehlo.subtract %5332, %5330 : tensor<1xf32> loc(#loc5260)
    %5334 = stablehlo.divide %5327, %5333 : tensor<1xf32> loc(#loc5253)
    %5335 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5336 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5337 = stablehlo.broadcast_in_dim %5336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5338 = stablehlo.add %5335, %5337 : tensor<1xf32> loc(#loc5259)
    %5339 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5340 = stablehlo.broadcast_in_dim %5339, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5341 = stablehlo.subtract %5338, %5340 : tensor<1xf32> loc(#loc5260)
    %5342 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5343 = stablehlo.broadcast_in_dim %5342, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5344 = stablehlo.power %5343, %5341 : tensor<1xf32> loc(#loc5261)
    %5345 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5346 = stablehlo.broadcast_in_dim %5345, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5347 = stablehlo.subtract %5346, %5344 : tensor<1xf32> loc(#loc5260)
    %5348 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5349 = stablehlo.broadcast_in_dim %5348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5350 = stablehlo.multiply %5349, %5347 : tensor<1xf32> loc(#loc5236)
    %5351 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5352 = stablehlo.broadcast_in_dim %5351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5353 = stablehlo.power %5352, %5338 : tensor<1xf32> loc(#loc5261)
    %5354 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5355 = stablehlo.broadcast_in_dim %5354, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5356 = stablehlo.subtract %5355, %5353 : tensor<1xf32> loc(#loc5260)
    %5357 = stablehlo.divide %5350, %5356 : tensor<1xf32> loc(#loc5253)
    %5358 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5359 = stablehlo.broadcast_in_dim %5358, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5360 = stablehlo.subtract %5359, %5334 : tensor<1xf32> loc(#loc5260)
    %5361 = stablehlo.broadcast_in_dim %5360, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5362 = stablehlo.broadcast_in_dim %5361, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5363 = stablehlo.multiply %5362, %4296 : tensor<1x2048xf32> loc(#loc5262)
    %5364 = stablehlo.broadcast_in_dim %5334, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5365 = stablehlo.broadcast_in_dim %5364, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5366 = stablehlo.multiply %5365, %arg204 : tensor<1x2048xf32> loc(#loc5262)
    %5367 = stablehlo.add %5363, %5366 : tensor<1x2048xf32> loc(#loc5263)
    %5368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5369 = stablehlo.broadcast_in_dim %5368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5370 = stablehlo.subtract %5369, %5357 : tensor<1xf32> loc(#loc5260)
    %5371 = stablehlo.multiply %4296, %4296 : tensor<1x2048xf32> loc(#loc5264)
    %5372 = stablehlo.broadcast_in_dim %5370, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5373 = stablehlo.broadcast_in_dim %5372, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5374 = stablehlo.multiply %5373, %5371 : tensor<1x2048xf32> loc(#loc5262)
    %5375 = stablehlo.broadcast_in_dim %5357, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5376 = stablehlo.broadcast_in_dim %5375, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5377 = stablehlo.multiply %5376, %arg380 : tensor<1x2048xf32> loc(#loc5262)
    %5378 = stablehlo.add %5374, %5377 : tensor<1x2048xf32> loc(#loc5263)
    %5379 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5380 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5381 = stablehlo.broadcast_in_dim %5380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5382 = stablehlo.add %5379, %5381 : tensor<1xf32> loc(#loc5259)
    %5383 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5384 = stablehlo.broadcast_in_dim %5383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5385 = stablehlo.subtract %5382, %5384 : tensor<1xf32> loc(#loc5260)
    %5386 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5387 = stablehlo.broadcast_in_dim %5386, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5388 = stablehlo.power %5387, %5385 : tensor<1xf32> loc(#loc5261)
    %5389 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5390 = stablehlo.broadcast_in_dim %5389, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5391 = stablehlo.subtract %5390, %5388 : tensor<1xf32> loc(#loc5260)
    %5392 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5393 = stablehlo.broadcast_in_dim %5392, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5394 = stablehlo.multiply %5393, %5391 : tensor<1xf32> loc(#loc5236)
    %5395 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5396 = stablehlo.broadcast_in_dim %5395, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5397 = stablehlo.power %5396, %5382 : tensor<1xf32> loc(#loc5261)
    %5398 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5399 = stablehlo.broadcast_in_dim %5398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5400 = stablehlo.subtract %5399, %5397 : tensor<1xf32> loc(#loc5260)
    %5401 = stablehlo.divide %5394, %5400 : tensor<1xf32> loc(#loc5253)
    %5402 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5403 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5404 = stablehlo.broadcast_in_dim %5403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5405 = stablehlo.add %5402, %5404 : tensor<1xf32> loc(#loc5259)
    %5406 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5407 = stablehlo.broadcast_in_dim %5406, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5408 = stablehlo.subtract %5405, %5407 : tensor<1xf32> loc(#loc5260)
    %5409 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5410 = stablehlo.broadcast_in_dim %5409, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5411 = stablehlo.power %5410, %5408 : tensor<1xf32> loc(#loc5261)
    %5412 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5413 = stablehlo.broadcast_in_dim %5412, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5414 = stablehlo.subtract %5413, %5411 : tensor<1xf32> loc(#loc5260)
    %5415 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5416 = stablehlo.broadcast_in_dim %5415, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5417 = stablehlo.multiply %5416, %5414 : tensor<1xf32> loc(#loc5236)
    %5418 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5419 = stablehlo.broadcast_in_dim %5418, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5420 = stablehlo.power %5419, %5405 : tensor<1xf32> loc(#loc5261)
    %5421 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5422 = stablehlo.broadcast_in_dim %5421, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5423 = stablehlo.subtract %5422, %5420 : tensor<1xf32> loc(#loc5260)
    %5424 = stablehlo.divide %5417, %5423 : tensor<1xf32> loc(#loc5253)
    %5425 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5426 = stablehlo.broadcast_in_dim %5425, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5427 = stablehlo.subtract %5426, %5401 : tensor<1xf32> loc(#loc5260)
    %5428 = stablehlo.broadcast_in_dim %5427, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5429 = stablehlo.broadcast_in_dim %5428, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5430 = stablehlo.multiply %5429, %4300 : tensor<1x2048xf32> loc(#loc5262)
    %5431 = stablehlo.broadcast_in_dim %5401, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5432 = stablehlo.broadcast_in_dim %5431, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5433 = stablehlo.multiply %5432, %arg205 : tensor<1x2048xf32> loc(#loc5262)
    %5434 = stablehlo.add %5430, %5433 : tensor<1x2048xf32> loc(#loc5263)
    %5435 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5436 = stablehlo.broadcast_in_dim %5435, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5437 = stablehlo.subtract %5436, %5424 : tensor<1xf32> loc(#loc5260)
    %5438 = stablehlo.multiply %4300, %4300 : tensor<1x2048xf32> loc(#loc5264)
    %5439 = stablehlo.broadcast_in_dim %5437, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5440 = stablehlo.broadcast_in_dim %5439, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5441 = stablehlo.multiply %5440, %5438 : tensor<1x2048xf32> loc(#loc5262)
    %5442 = stablehlo.broadcast_in_dim %5424, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5443 = stablehlo.broadcast_in_dim %5442, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5444 = stablehlo.multiply %5443, %arg381 : tensor<1x2048xf32> loc(#loc5262)
    %5445 = stablehlo.add %5441, %5444 : tensor<1x2048xf32> loc(#loc5263)
    %5446 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5447 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5448 = stablehlo.broadcast_in_dim %5447, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5449 = stablehlo.add %5446, %5448 : tensor<1xf32> loc(#loc5259)
    %5450 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5451 = stablehlo.broadcast_in_dim %5450, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5452 = stablehlo.subtract %5449, %5451 : tensor<1xf32> loc(#loc5260)
    %5453 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5454 = stablehlo.broadcast_in_dim %5453, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5455 = stablehlo.power %5454, %5452 : tensor<1xf32> loc(#loc5261)
    %5456 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5457 = stablehlo.broadcast_in_dim %5456, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5458 = stablehlo.subtract %5457, %5455 : tensor<1xf32> loc(#loc5260)
    %5459 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5460 = stablehlo.broadcast_in_dim %5459, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5461 = stablehlo.multiply %5460, %5458 : tensor<1xf32> loc(#loc5236)
    %5462 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5463 = stablehlo.broadcast_in_dim %5462, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5464 = stablehlo.power %5463, %5449 : tensor<1xf32> loc(#loc5261)
    %5465 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5466 = stablehlo.broadcast_in_dim %5465, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5467 = stablehlo.subtract %5466, %5464 : tensor<1xf32> loc(#loc5260)
    %5468 = stablehlo.divide %5461, %5467 : tensor<1xf32> loc(#loc5253)
    %5469 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5471 = stablehlo.broadcast_in_dim %5470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5472 = stablehlo.add %5469, %5471 : tensor<1xf32> loc(#loc5259)
    %5473 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5474 = stablehlo.broadcast_in_dim %5473, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5475 = stablehlo.subtract %5472, %5474 : tensor<1xf32> loc(#loc5260)
    %5476 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5477 = stablehlo.broadcast_in_dim %5476, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5478 = stablehlo.power %5477, %5475 : tensor<1xf32> loc(#loc5261)
    %5479 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5480 = stablehlo.broadcast_in_dim %5479, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5481 = stablehlo.subtract %5480, %5478 : tensor<1xf32> loc(#loc5260)
    %5482 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5483 = stablehlo.broadcast_in_dim %5482, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5484 = stablehlo.multiply %5483, %5481 : tensor<1xf32> loc(#loc5236)
    %5485 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5486 = stablehlo.broadcast_in_dim %5485, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5487 = stablehlo.power %5486, %5472 : tensor<1xf32> loc(#loc5261)
    %5488 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5489 = stablehlo.broadcast_in_dim %5488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5490 = stablehlo.subtract %5489, %5487 : tensor<1xf32> loc(#loc5260)
    %5491 = stablehlo.divide %5484, %5490 : tensor<1xf32> loc(#loc5253)
    %5492 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5493 = stablehlo.broadcast_in_dim %5492, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5494 = stablehlo.subtract %5493, %5468 : tensor<1xf32> loc(#loc5260)
    %5495 = stablehlo.broadcast_in_dim %5494, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5496 = stablehlo.broadcast_in_dim %5495, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5497 = stablehlo.multiply %5496, %4304 : tensor<1x2048xf32> loc(#loc5262)
    %5498 = stablehlo.broadcast_in_dim %5468, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5499 = stablehlo.broadcast_in_dim %5498, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5500 = stablehlo.multiply %5499, %arg206 : tensor<1x2048xf32> loc(#loc5262)
    %5501 = stablehlo.add %5497, %5500 : tensor<1x2048xf32> loc(#loc5263)
    %5502 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5503 = stablehlo.broadcast_in_dim %5502, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5504 = stablehlo.subtract %5503, %5491 : tensor<1xf32> loc(#loc5260)
    %5505 = stablehlo.multiply %4304, %4304 : tensor<1x2048xf32> loc(#loc5264)
    %5506 = stablehlo.broadcast_in_dim %5504, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5507 = stablehlo.broadcast_in_dim %5506, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5508 = stablehlo.multiply %5507, %5505 : tensor<1x2048xf32> loc(#loc5262)
    %5509 = stablehlo.broadcast_in_dim %5491, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5510 = stablehlo.broadcast_in_dim %5509, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5511 = stablehlo.multiply %5510, %arg382 : tensor<1x2048xf32> loc(#loc5262)
    %5512 = stablehlo.add %5508, %5511 : tensor<1x2048xf32> loc(#loc5263)
    %5513 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5514 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5515 = stablehlo.broadcast_in_dim %5514, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5516 = stablehlo.add %5513, %5515 : tensor<1xf32> loc(#loc5259)
    %5517 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5518 = stablehlo.broadcast_in_dim %5517, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5519 = stablehlo.subtract %5516, %5518 : tensor<1xf32> loc(#loc5260)
    %5520 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5521 = stablehlo.broadcast_in_dim %5520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5522 = stablehlo.power %5521, %5519 : tensor<1xf32> loc(#loc5261)
    %5523 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5524 = stablehlo.broadcast_in_dim %5523, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5525 = stablehlo.subtract %5524, %5522 : tensor<1xf32> loc(#loc5260)
    %5526 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5527 = stablehlo.broadcast_in_dim %5526, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5528 = stablehlo.multiply %5527, %5525 : tensor<1xf32> loc(#loc5236)
    %5529 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5530 = stablehlo.broadcast_in_dim %5529, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5531 = stablehlo.power %5530, %5516 : tensor<1xf32> loc(#loc5261)
    %5532 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5533 = stablehlo.broadcast_in_dim %5532, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5534 = stablehlo.subtract %5533, %5531 : tensor<1xf32> loc(#loc5260)
    %5535 = stablehlo.divide %5528, %5534 : tensor<1xf32> loc(#loc5253)
    %5536 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5537 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5538 = stablehlo.broadcast_in_dim %5537, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5539 = stablehlo.add %5536, %5538 : tensor<1xf32> loc(#loc5259)
    %5540 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5541 = stablehlo.broadcast_in_dim %5540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5542 = stablehlo.subtract %5539, %5541 : tensor<1xf32> loc(#loc5260)
    %5543 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5544 = stablehlo.broadcast_in_dim %5543, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5545 = stablehlo.power %5544, %5542 : tensor<1xf32> loc(#loc5261)
    %5546 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5547 = stablehlo.broadcast_in_dim %5546, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5548 = stablehlo.subtract %5547, %5545 : tensor<1xf32> loc(#loc5260)
    %5549 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5550 = stablehlo.broadcast_in_dim %5549, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5551 = stablehlo.multiply %5550, %5548 : tensor<1xf32> loc(#loc5236)
    %5552 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5553 = stablehlo.broadcast_in_dim %5552, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5554 = stablehlo.power %5553, %5539 : tensor<1xf32> loc(#loc5261)
    %5555 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5556 = stablehlo.broadcast_in_dim %5555, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5557 = stablehlo.subtract %5556, %5554 : tensor<1xf32> loc(#loc5260)
    %5558 = stablehlo.divide %5551, %5557 : tensor<1xf32> loc(#loc5253)
    %5559 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5560 = stablehlo.broadcast_in_dim %5559, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5561 = stablehlo.subtract %5560, %5535 : tensor<1xf32> loc(#loc5260)
    %5562 = stablehlo.broadcast_in_dim %5561, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %5563 = stablehlo.broadcast_in_dim %5562, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5564 = stablehlo.multiply %5563, %4308 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5565 = stablehlo.broadcast_in_dim %5535, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %5566 = stablehlo.broadcast_in_dim %5565, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5567 = stablehlo.multiply %5566, %arg207 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5568 = stablehlo.add %5564, %5567 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %5569 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5570 = stablehlo.broadcast_in_dim %5569, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5571 = stablehlo.subtract %5570, %5558 : tensor<1xf32> loc(#loc5260)
    %5572 = stablehlo.multiply %4308, %4308 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %5573 = stablehlo.broadcast_in_dim %5571, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %5574 = stablehlo.broadcast_in_dim %5573, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5575 = stablehlo.multiply %5574, %5572 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5576 = stablehlo.broadcast_in_dim %5558, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %5577 = stablehlo.broadcast_in_dim %5576, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5578 = stablehlo.multiply %5577, %arg383 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %5579 = stablehlo.add %5575, %5578 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %5580 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5581 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5582 = stablehlo.broadcast_in_dim %5581, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5583 = stablehlo.add %5580, %5582 : tensor<1xf32> loc(#loc5259)
    %5584 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5585 = stablehlo.broadcast_in_dim %5584, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5586 = stablehlo.subtract %5583, %5585 : tensor<1xf32> loc(#loc5260)
    %5587 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5588 = stablehlo.broadcast_in_dim %5587, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5589 = stablehlo.power %5588, %5586 : tensor<1xf32> loc(#loc5261)
    %5590 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5591 = stablehlo.broadcast_in_dim %5590, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5592 = stablehlo.subtract %5591, %5589 : tensor<1xf32> loc(#loc5260)
    %5593 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5594 = stablehlo.broadcast_in_dim %5593, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5595 = stablehlo.multiply %5594, %5592 : tensor<1xf32> loc(#loc5236)
    %5596 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5597 = stablehlo.broadcast_in_dim %5596, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5598 = stablehlo.power %5597, %5583 : tensor<1xf32> loc(#loc5261)
    %5599 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5600 = stablehlo.broadcast_in_dim %5599, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5601 = stablehlo.subtract %5600, %5598 : tensor<1xf32> loc(#loc5260)
    %5602 = stablehlo.divide %5595, %5601 : tensor<1xf32> loc(#loc5253)
    %5603 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5604 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5605 = stablehlo.broadcast_in_dim %5604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5606 = stablehlo.add %5603, %5605 : tensor<1xf32> loc(#loc5259)
    %5607 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5608 = stablehlo.broadcast_in_dim %5607, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5609 = stablehlo.subtract %5606, %5608 : tensor<1xf32> loc(#loc5260)
    %5610 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5611 = stablehlo.broadcast_in_dim %5610, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5612 = stablehlo.power %5611, %5609 : tensor<1xf32> loc(#loc5261)
    %5613 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5614 = stablehlo.broadcast_in_dim %5613, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5615 = stablehlo.subtract %5614, %5612 : tensor<1xf32> loc(#loc5260)
    %5616 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5617 = stablehlo.broadcast_in_dim %5616, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5618 = stablehlo.multiply %5617, %5615 : tensor<1xf32> loc(#loc5236)
    %5619 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5620 = stablehlo.broadcast_in_dim %5619, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5621 = stablehlo.power %5620, %5606 : tensor<1xf32> loc(#loc5261)
    %5622 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5623 = stablehlo.broadcast_in_dim %5622, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5624 = stablehlo.subtract %5623, %5621 : tensor<1xf32> loc(#loc5260)
    %5625 = stablehlo.divide %5618, %5624 : tensor<1xf32> loc(#loc5253)
    %5626 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5627 = stablehlo.broadcast_in_dim %5626, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5628 = stablehlo.subtract %5627, %5602 : tensor<1xf32> loc(#loc5260)
    %5629 = stablehlo.broadcast_in_dim %5628, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5630 = stablehlo.broadcast_in_dim %5629, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %5631 = stablehlo.multiply %5630, %4312 : tensor<1x64xf32> loc(#loc5262)
    %5632 = stablehlo.broadcast_in_dim %5602, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5633 = stablehlo.broadcast_in_dim %5632, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %5634 = stablehlo.multiply %5633, %arg208 : tensor<1x64xf32> loc(#loc5262)
    %5635 = stablehlo.add %5631, %5634 : tensor<1x64xf32> loc(#loc5263)
    %5636 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5637 = stablehlo.broadcast_in_dim %5636, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5638 = stablehlo.subtract %5637, %5625 : tensor<1xf32> loc(#loc5260)
    %5639 = stablehlo.multiply %4312, %4312 : tensor<1x64xf32> loc(#loc5264)
    %5640 = stablehlo.broadcast_in_dim %5638, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5641 = stablehlo.broadcast_in_dim %5640, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %5642 = stablehlo.multiply %5641, %5639 : tensor<1x64xf32> loc(#loc5262)
    %5643 = stablehlo.broadcast_in_dim %5625, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5644 = stablehlo.broadcast_in_dim %5643, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %5645 = stablehlo.multiply %5644, %arg384 : tensor<1x64xf32> loc(#loc5262)
    %5646 = stablehlo.add %5642, %5645 : tensor<1x64xf32> loc(#loc5263)
    %5647 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5648 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5649 = stablehlo.broadcast_in_dim %5648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5650 = stablehlo.add %5647, %5649 : tensor<1xf32> loc(#loc5259)
    %5651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5652 = stablehlo.broadcast_in_dim %5651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5653 = stablehlo.subtract %5650, %5652 : tensor<1xf32> loc(#loc5260)
    %5654 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5655 = stablehlo.broadcast_in_dim %5654, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5656 = stablehlo.power %5655, %5653 : tensor<1xf32> loc(#loc5261)
    %5657 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5658 = stablehlo.broadcast_in_dim %5657, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5659 = stablehlo.subtract %5658, %5656 : tensor<1xf32> loc(#loc5260)
    %5660 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5661 = stablehlo.broadcast_in_dim %5660, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5662 = stablehlo.multiply %5661, %5659 : tensor<1xf32> loc(#loc5236)
    %5663 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5664 = stablehlo.broadcast_in_dim %5663, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5665 = stablehlo.power %5664, %5650 : tensor<1xf32> loc(#loc5261)
    %5666 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5667 = stablehlo.broadcast_in_dim %5666, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5668 = stablehlo.subtract %5667, %5665 : tensor<1xf32> loc(#loc5260)
    %5669 = stablehlo.divide %5662, %5668 : tensor<1xf32> loc(#loc5253)
    %5670 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5672 = stablehlo.broadcast_in_dim %5671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5673 = stablehlo.add %5670, %5672 : tensor<1xf32> loc(#loc5259)
    %5674 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5675 = stablehlo.broadcast_in_dim %5674, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5676 = stablehlo.subtract %5673, %5675 : tensor<1xf32> loc(#loc5260)
    %5677 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5678 = stablehlo.broadcast_in_dim %5677, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5679 = stablehlo.power %5678, %5676 : tensor<1xf32> loc(#loc5261)
    %5680 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5681 = stablehlo.broadcast_in_dim %5680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5682 = stablehlo.subtract %5681, %5679 : tensor<1xf32> loc(#loc5260)
    %5683 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5684 = stablehlo.broadcast_in_dim %5683, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5685 = stablehlo.multiply %5684, %5682 : tensor<1xf32> loc(#loc5236)
    %5686 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5687 = stablehlo.broadcast_in_dim %5686, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5688 = stablehlo.power %5687, %5673 : tensor<1xf32> loc(#loc5261)
    %5689 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5690 = stablehlo.broadcast_in_dim %5689, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5691 = stablehlo.subtract %5690, %5688 : tensor<1xf32> loc(#loc5260)
    %5692 = stablehlo.divide %5685, %5691 : tensor<1xf32> loc(#loc5253)
    %5693 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5694 = stablehlo.broadcast_in_dim %5693, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5695 = stablehlo.subtract %5694, %5669 : tensor<1xf32> loc(#loc5260)
    %5696 = stablehlo.broadcast_in_dim %5695, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %5697 = stablehlo.broadcast_in_dim %5696, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %5698 = stablehlo.multiply %5697, %4316 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %5699 = stablehlo.broadcast_in_dim %5669, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %5700 = stablehlo.broadcast_in_dim %5699, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %5701 = stablehlo.multiply %5700, %arg209 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %5702 = stablehlo.add %5698, %5701 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %5703 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5704 = stablehlo.broadcast_in_dim %5703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5705 = stablehlo.subtract %5704, %5692 : tensor<1xf32> loc(#loc5260)
    %5706 = stablehlo.multiply %4316, %4316 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %5707 = stablehlo.broadcast_in_dim %5705, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %5708 = stablehlo.broadcast_in_dim %5707, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %5709 = stablehlo.multiply %5708, %5706 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %5710 = stablehlo.broadcast_in_dim %5692, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %5711 = stablehlo.broadcast_in_dim %5710, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %5712 = stablehlo.multiply %5711, %arg385 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %5713 = stablehlo.add %5709, %5712 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %5714 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5715 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5716 = stablehlo.broadcast_in_dim %5715, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5717 = stablehlo.add %5714, %5716 : tensor<1xf32> loc(#loc5259)
    %5718 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5719 = stablehlo.broadcast_in_dim %5718, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5720 = stablehlo.subtract %5717, %5719 : tensor<1xf32> loc(#loc5260)
    %5721 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5722 = stablehlo.broadcast_in_dim %5721, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5723 = stablehlo.power %5722, %5720 : tensor<1xf32> loc(#loc5261)
    %5724 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5725 = stablehlo.broadcast_in_dim %5724, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5726 = stablehlo.subtract %5725, %5723 : tensor<1xf32> loc(#loc5260)
    %5727 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5728 = stablehlo.broadcast_in_dim %5727, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5729 = stablehlo.multiply %5728, %5726 : tensor<1xf32> loc(#loc5236)
    %5730 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5731 = stablehlo.broadcast_in_dim %5730, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5732 = stablehlo.power %5731, %5717 : tensor<1xf32> loc(#loc5261)
    %5733 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5734 = stablehlo.broadcast_in_dim %5733, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5735 = stablehlo.subtract %5734, %5732 : tensor<1xf32> loc(#loc5260)
    %5736 = stablehlo.divide %5729, %5735 : tensor<1xf32> loc(#loc5253)
    %5737 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5738 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5739 = stablehlo.broadcast_in_dim %5738, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5740 = stablehlo.add %5737, %5739 : tensor<1xf32> loc(#loc5259)
    %5741 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5742 = stablehlo.broadcast_in_dim %5741, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5743 = stablehlo.subtract %5740, %5742 : tensor<1xf32> loc(#loc5260)
    %5744 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5745 = stablehlo.broadcast_in_dim %5744, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5746 = stablehlo.power %5745, %5743 : tensor<1xf32> loc(#loc5261)
    %5747 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5748 = stablehlo.broadcast_in_dim %5747, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5749 = stablehlo.subtract %5748, %5746 : tensor<1xf32> loc(#loc5260)
    %5750 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5751 = stablehlo.broadcast_in_dim %5750, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5752 = stablehlo.multiply %5751, %5749 : tensor<1xf32> loc(#loc5236)
    %5753 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5754 = stablehlo.broadcast_in_dim %5753, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5755 = stablehlo.power %5754, %5740 : tensor<1xf32> loc(#loc5261)
    %5756 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5757 = stablehlo.broadcast_in_dim %5756, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5758 = stablehlo.subtract %5757, %5755 : tensor<1xf32> loc(#loc5260)
    %5759 = stablehlo.divide %5752, %5758 : tensor<1xf32> loc(#loc5253)
    %5760 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5761 = stablehlo.broadcast_in_dim %5760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5762 = stablehlo.subtract %5761, %5736 : tensor<1xf32> loc(#loc5260)
    %5763 = stablehlo.broadcast_in_dim %5762, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5764 = stablehlo.broadcast_in_dim %5763, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5765 = stablehlo.multiply %5764, %4320 : tensor<1x8192xf32> loc(#loc5262)
    %5766 = stablehlo.broadcast_in_dim %5736, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5767 = stablehlo.broadcast_in_dim %5766, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5768 = stablehlo.multiply %5767, %arg210 : tensor<1x8192xf32> loc(#loc5262)
    %5769 = stablehlo.add %5765, %5768 : tensor<1x8192xf32> loc(#loc5263)
    %5770 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5771 = stablehlo.broadcast_in_dim %5770, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5772 = stablehlo.subtract %5771, %5759 : tensor<1xf32> loc(#loc5260)
    %5773 = stablehlo.multiply %4320, %4320 : tensor<1x8192xf32> loc(#loc5264)
    %5774 = stablehlo.broadcast_in_dim %5772, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5775 = stablehlo.broadcast_in_dim %5774, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5776 = stablehlo.multiply %5775, %5773 : tensor<1x8192xf32> loc(#loc5262)
    %5777 = stablehlo.broadcast_in_dim %5759, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5778 = stablehlo.broadcast_in_dim %5777, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %5779 = stablehlo.multiply %5778, %arg386 : tensor<1x8192xf32> loc(#loc5262)
    %5780 = stablehlo.add %5776, %5779 : tensor<1x8192xf32> loc(#loc5263)
    %5781 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5782 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5783 = stablehlo.broadcast_in_dim %5782, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5784 = stablehlo.add %5781, %5783 : tensor<1xf32> loc(#loc5259)
    %5785 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5786 = stablehlo.broadcast_in_dim %5785, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5787 = stablehlo.subtract %5784, %5786 : tensor<1xf32> loc(#loc5260)
    %5788 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5789 = stablehlo.broadcast_in_dim %5788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5790 = stablehlo.power %5789, %5787 : tensor<1xf32> loc(#loc5261)
    %5791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5792 = stablehlo.broadcast_in_dim %5791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5793 = stablehlo.subtract %5792, %5790 : tensor<1xf32> loc(#loc5260)
    %5794 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5795 = stablehlo.broadcast_in_dim %5794, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5796 = stablehlo.multiply %5795, %5793 : tensor<1xf32> loc(#loc5236)
    %5797 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5798 = stablehlo.broadcast_in_dim %5797, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5799 = stablehlo.power %5798, %5784 : tensor<1xf32> loc(#loc5261)
    %5800 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5801 = stablehlo.broadcast_in_dim %5800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5802 = stablehlo.subtract %5801, %5799 : tensor<1xf32> loc(#loc5260)
    %5803 = stablehlo.divide %5796, %5802 : tensor<1xf32> loc(#loc5253)
    %5804 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5805 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5806 = stablehlo.broadcast_in_dim %5805, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5807 = stablehlo.add %5804, %5806 : tensor<1xf32> loc(#loc5259)
    %5808 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5809 = stablehlo.broadcast_in_dim %5808, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5810 = stablehlo.subtract %5807, %5809 : tensor<1xf32> loc(#loc5260)
    %5811 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5812 = stablehlo.broadcast_in_dim %5811, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5813 = stablehlo.power %5812, %5810 : tensor<1xf32> loc(#loc5261)
    %5814 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5815 = stablehlo.broadcast_in_dim %5814, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5816 = stablehlo.subtract %5815, %5813 : tensor<1xf32> loc(#loc5260)
    %5817 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5818 = stablehlo.broadcast_in_dim %5817, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5819 = stablehlo.multiply %5818, %5816 : tensor<1xf32> loc(#loc5236)
    %5820 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5821 = stablehlo.broadcast_in_dim %5820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5822 = stablehlo.power %5821, %5807 : tensor<1xf32> loc(#loc5261)
    %5823 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5824 = stablehlo.broadcast_in_dim %5823, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5825 = stablehlo.subtract %5824, %5822 : tensor<1xf32> loc(#loc5260)
    %5826 = stablehlo.divide %5819, %5825 : tensor<1xf32> loc(#loc5253)
    %5827 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5828 = stablehlo.broadcast_in_dim %5827, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5829 = stablehlo.subtract %5828, %5803 : tensor<1xf32> loc(#loc5260)
    %5830 = stablehlo.broadcast_in_dim %5829, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5831 = stablehlo.broadcast_in_dim %5830, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5832 = stablehlo.multiply %5831, %4324 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5833 = stablehlo.broadcast_in_dim %5803, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5834 = stablehlo.broadcast_in_dim %5833, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5835 = stablehlo.multiply %5834, %arg211 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5836 = stablehlo.add %5832, %5835 : tensor<1x2048x8192xf32> loc(#loc5263)
    %5837 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5838 = stablehlo.broadcast_in_dim %5837, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5839 = stablehlo.subtract %5838, %5826 : tensor<1xf32> loc(#loc5260)
    %5840 = stablehlo.multiply %4324, %4324 : tensor<1x2048x8192xf32> loc(#loc5264)
    %5841 = stablehlo.broadcast_in_dim %5839, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5842 = stablehlo.broadcast_in_dim %5841, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5843 = stablehlo.multiply %5842, %5840 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5844 = stablehlo.broadcast_in_dim %5826, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5845 = stablehlo.broadcast_in_dim %5844, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %5846 = stablehlo.multiply %5845, %arg387 : tensor<1x2048x8192xf32> loc(#loc5262)
    %5847 = stablehlo.add %5843, %5846 : tensor<1x2048x8192xf32> loc(#loc5263)
    %5848 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5849 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5850 = stablehlo.broadcast_in_dim %5849, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5851 = stablehlo.add %5848, %5850 : tensor<1xf32> loc(#loc5259)
    %5852 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5853 = stablehlo.broadcast_in_dim %5852, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5854 = stablehlo.subtract %5851, %5853 : tensor<1xf32> loc(#loc5260)
    %5855 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5856 = stablehlo.broadcast_in_dim %5855, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5857 = stablehlo.power %5856, %5854 : tensor<1xf32> loc(#loc5261)
    %5858 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5859 = stablehlo.broadcast_in_dim %5858, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5860 = stablehlo.subtract %5859, %5857 : tensor<1xf32> loc(#loc5260)
    %5861 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5862 = stablehlo.broadcast_in_dim %5861, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5863 = stablehlo.multiply %5862, %5860 : tensor<1xf32> loc(#loc5236)
    %5864 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5865 = stablehlo.broadcast_in_dim %5864, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5866 = stablehlo.power %5865, %5851 : tensor<1xf32> loc(#loc5261)
    %5867 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5868 = stablehlo.broadcast_in_dim %5867, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5869 = stablehlo.subtract %5868, %5866 : tensor<1xf32> loc(#loc5260)
    %5870 = stablehlo.divide %5863, %5869 : tensor<1xf32> loc(#loc5253)
    %5871 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5872 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5873 = stablehlo.broadcast_in_dim %5872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5874 = stablehlo.add %5871, %5873 : tensor<1xf32> loc(#loc5259)
    %5875 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5876 = stablehlo.broadcast_in_dim %5875, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5877 = stablehlo.subtract %5874, %5876 : tensor<1xf32> loc(#loc5260)
    %5878 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5879 = stablehlo.broadcast_in_dim %5878, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5880 = stablehlo.power %5879, %5877 : tensor<1xf32> loc(#loc5261)
    %5881 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5882 = stablehlo.broadcast_in_dim %5881, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5883 = stablehlo.subtract %5882, %5880 : tensor<1xf32> loc(#loc5260)
    %5884 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5885 = stablehlo.broadcast_in_dim %5884, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5886 = stablehlo.multiply %5885, %5883 : tensor<1xf32> loc(#loc5236)
    %5887 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5888 = stablehlo.broadcast_in_dim %5887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5889 = stablehlo.power %5888, %5874 : tensor<1xf32> loc(#loc5261)
    %5890 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5891 = stablehlo.broadcast_in_dim %5890, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5892 = stablehlo.subtract %5891, %5889 : tensor<1xf32> loc(#loc5260)
    %5893 = stablehlo.divide %5886, %5892 : tensor<1xf32> loc(#loc5253)
    %5894 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5895 = stablehlo.broadcast_in_dim %5894, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5896 = stablehlo.subtract %5895, %5870 : tensor<1xf32> loc(#loc5260)
    %5897 = stablehlo.broadcast_in_dim %5896, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5898 = stablehlo.broadcast_in_dim %5897, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5899 = stablehlo.multiply %5898, %4328 : tensor<1x2048xf32> loc(#loc5262)
    %5900 = stablehlo.broadcast_in_dim %5870, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5901 = stablehlo.broadcast_in_dim %5900, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5902 = stablehlo.multiply %5901, %arg212 : tensor<1x2048xf32> loc(#loc5262)
    %5903 = stablehlo.add %5899, %5902 : tensor<1x2048xf32> loc(#loc5263)
    %5904 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5905 = stablehlo.broadcast_in_dim %5904, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5906 = stablehlo.subtract %5905, %5893 : tensor<1xf32> loc(#loc5260)
    %5907 = stablehlo.multiply %4328, %4328 : tensor<1x2048xf32> loc(#loc5264)
    %5908 = stablehlo.broadcast_in_dim %5906, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5909 = stablehlo.broadcast_in_dim %5908, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5910 = stablehlo.multiply %5909, %5907 : tensor<1x2048xf32> loc(#loc5262)
    %5911 = stablehlo.broadcast_in_dim %5893, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %5912 = stablehlo.broadcast_in_dim %5911, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %5913 = stablehlo.multiply %5912, %arg388 : tensor<1x2048xf32> loc(#loc5262)
    %5914 = stablehlo.add %5910, %5913 : tensor<1x2048xf32> loc(#loc5263)
    %5915 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5916 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5917 = stablehlo.broadcast_in_dim %5916, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5918 = stablehlo.add %5915, %5917 : tensor<1xf32> loc(#loc5259)
    %5919 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5920 = stablehlo.broadcast_in_dim %5919, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5921 = stablehlo.subtract %5918, %5920 : tensor<1xf32> loc(#loc5260)
    %5922 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5923 = stablehlo.broadcast_in_dim %5922, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5924 = stablehlo.power %5923, %5921 : tensor<1xf32> loc(#loc5261)
    %5925 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5926 = stablehlo.broadcast_in_dim %5925, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5927 = stablehlo.subtract %5926, %5924 : tensor<1xf32> loc(#loc5260)
    %5928 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5929 = stablehlo.broadcast_in_dim %5928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5930 = stablehlo.multiply %5929, %5927 : tensor<1xf32> loc(#loc5236)
    %5931 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5932 = stablehlo.broadcast_in_dim %5931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5933 = stablehlo.power %5932, %5918 : tensor<1xf32> loc(#loc5261)
    %5934 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5935 = stablehlo.broadcast_in_dim %5934, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5936 = stablehlo.subtract %5935, %5933 : tensor<1xf32> loc(#loc5260)
    %5937 = stablehlo.divide %5930, %5936 : tensor<1xf32> loc(#loc5253)
    %5938 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5939 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5940 = stablehlo.broadcast_in_dim %5939, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5941 = stablehlo.add %5938, %5940 : tensor<1xf32> loc(#loc5259)
    %5942 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5943 = stablehlo.broadcast_in_dim %5942, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5944 = stablehlo.subtract %5941, %5943 : tensor<1xf32> loc(#loc5260)
    %5945 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5946 = stablehlo.broadcast_in_dim %5945, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5947 = stablehlo.power %5946, %5944 : tensor<1xf32> loc(#loc5261)
    %5948 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5949 = stablehlo.broadcast_in_dim %5948, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5950 = stablehlo.subtract %5949, %5947 : tensor<1xf32> loc(#loc5260)
    %5951 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5952 = stablehlo.broadcast_in_dim %5951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5953 = stablehlo.multiply %5952, %5950 : tensor<1xf32> loc(#loc5236)
    %5954 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %5955 = stablehlo.broadcast_in_dim %5954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5956 = stablehlo.power %5955, %5941 : tensor<1xf32> loc(#loc5261)
    %5957 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5958 = stablehlo.broadcast_in_dim %5957, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5959 = stablehlo.subtract %5958, %5956 : tensor<1xf32> loc(#loc5260)
    %5960 = stablehlo.divide %5953, %5959 : tensor<1xf32> loc(#loc5253)
    %5961 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5962 = stablehlo.broadcast_in_dim %5961, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5963 = stablehlo.subtract %5962, %5937 : tensor<1xf32> loc(#loc5260)
    %5964 = stablehlo.broadcast_in_dim %5963, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5965 = stablehlo.broadcast_in_dim %5964, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5966 = stablehlo.multiply %5965, %4332 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5967 = stablehlo.broadcast_in_dim %5937, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5968 = stablehlo.broadcast_in_dim %5967, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5969 = stablehlo.multiply %5968, %arg213 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5970 = stablehlo.add %5966, %5969 : tensor<1x8192x2048xf32> loc(#loc5263)
    %5971 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5972 = stablehlo.broadcast_in_dim %5971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5973 = stablehlo.subtract %5972, %5960 : tensor<1xf32> loc(#loc5260)
    %5974 = stablehlo.multiply %4332, %4332 : tensor<1x8192x2048xf32> loc(#loc5264)
    %5975 = stablehlo.broadcast_in_dim %5973, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5976 = stablehlo.broadcast_in_dim %5975, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5977 = stablehlo.multiply %5976, %5974 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5978 = stablehlo.broadcast_in_dim %5960, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %5979 = stablehlo.broadcast_in_dim %5978, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %5980 = stablehlo.multiply %5979, %arg389 : tensor<1x8192x2048xf32> loc(#loc5262)
    %5981 = stablehlo.add %5977, %5980 : tensor<1x8192x2048xf32> loc(#loc5263)
    %5982 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %5983 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5984 = stablehlo.broadcast_in_dim %5983, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %5985 = stablehlo.add %5982, %5984 : tensor<1xf32> loc(#loc5259)
    %5986 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5987 = stablehlo.broadcast_in_dim %5986, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5988 = stablehlo.subtract %5985, %5987 : tensor<1xf32> loc(#loc5260)
    %5989 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5990 = stablehlo.broadcast_in_dim %5989, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %5991 = stablehlo.power %5990, %5988 : tensor<1xf32> loc(#loc5261)
    %5992 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %5993 = stablehlo.broadcast_in_dim %5992, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %5994 = stablehlo.subtract %5993, %5991 : tensor<1xf32> loc(#loc5260)
    %5995 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5996 = stablehlo.broadcast_in_dim %5995, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %5997 = stablehlo.multiply %5996, %5994 : tensor<1xf32> loc(#loc5236)
    %5998 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %5999 = stablehlo.broadcast_in_dim %5998, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6000 = stablehlo.power %5999, %5985 : tensor<1xf32> loc(#loc5261)
    %6001 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6002 = stablehlo.broadcast_in_dim %6001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6003 = stablehlo.subtract %6002, %6000 : tensor<1xf32> loc(#loc5260)
    %6004 = stablehlo.divide %5997, %6003 : tensor<1xf32> loc(#loc5253)
    %6005 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6006 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6007 = stablehlo.broadcast_in_dim %6006, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6008 = stablehlo.add %6005, %6007 : tensor<1xf32> loc(#loc5259)
    %6009 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6010 = stablehlo.broadcast_in_dim %6009, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6011 = stablehlo.subtract %6008, %6010 : tensor<1xf32> loc(#loc5260)
    %6012 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6013 = stablehlo.broadcast_in_dim %6012, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6014 = stablehlo.power %6013, %6011 : tensor<1xf32> loc(#loc5261)
    %6015 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6016 = stablehlo.broadcast_in_dim %6015, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6017 = stablehlo.subtract %6016, %6014 : tensor<1xf32> loc(#loc5260)
    %6018 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6019 = stablehlo.broadcast_in_dim %6018, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6020 = stablehlo.multiply %6019, %6017 : tensor<1xf32> loc(#loc5236)
    %6021 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6022 = stablehlo.broadcast_in_dim %6021, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6023 = stablehlo.power %6022, %6008 : tensor<1xf32> loc(#loc5261)
    %6024 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6025 = stablehlo.broadcast_in_dim %6024, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6026 = stablehlo.subtract %6025, %6023 : tensor<1xf32> loc(#loc5260)
    %6027 = stablehlo.divide %6020, %6026 : tensor<1xf32> loc(#loc5253)
    %6028 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6029 = stablehlo.broadcast_in_dim %6028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6030 = stablehlo.subtract %6029, %6004 : tensor<1xf32> loc(#loc5260)
    %6031 = stablehlo.broadcast_in_dim %6030, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6032 = stablehlo.broadcast_in_dim %6031, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6033 = stablehlo.multiply %6032, %4336 : tensor<1x2048xf32> loc(#loc5262)
    %6034 = stablehlo.broadcast_in_dim %6004, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6035 = stablehlo.broadcast_in_dim %6034, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6036 = stablehlo.multiply %6035, %arg214 : tensor<1x2048xf32> loc(#loc5262)
    %6037 = stablehlo.add %6033, %6036 : tensor<1x2048xf32> loc(#loc5263)
    %6038 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6039 = stablehlo.broadcast_in_dim %6038, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6040 = stablehlo.subtract %6039, %6027 : tensor<1xf32> loc(#loc5260)
    %6041 = stablehlo.multiply %4336, %4336 : tensor<1x2048xf32> loc(#loc5264)
    %6042 = stablehlo.broadcast_in_dim %6040, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6043 = stablehlo.broadcast_in_dim %6042, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6044 = stablehlo.multiply %6043, %6041 : tensor<1x2048xf32> loc(#loc5262)
    %6045 = stablehlo.broadcast_in_dim %6027, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6046 = stablehlo.broadcast_in_dim %6045, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6047 = stablehlo.multiply %6046, %arg390 : tensor<1x2048xf32> loc(#loc5262)
    %6048 = stablehlo.add %6044, %6047 : tensor<1x2048xf32> loc(#loc5263)
    %6049 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6050 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6051 = stablehlo.broadcast_in_dim %6050, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6052 = stablehlo.add %6049, %6051 : tensor<1xf32> loc(#loc5259)
    %6053 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6054 = stablehlo.broadcast_in_dim %6053, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6055 = stablehlo.subtract %6052, %6054 : tensor<1xf32> loc(#loc5260)
    %6056 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6057 = stablehlo.broadcast_in_dim %6056, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6058 = stablehlo.power %6057, %6055 : tensor<1xf32> loc(#loc5261)
    %6059 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6060 = stablehlo.broadcast_in_dim %6059, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6061 = stablehlo.subtract %6060, %6058 : tensor<1xf32> loc(#loc5260)
    %6062 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6063 = stablehlo.broadcast_in_dim %6062, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6064 = stablehlo.multiply %6063, %6061 : tensor<1xf32> loc(#loc5236)
    %6065 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6066 = stablehlo.broadcast_in_dim %6065, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6067 = stablehlo.power %6066, %6052 : tensor<1xf32> loc(#loc5261)
    %6068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6069 = stablehlo.broadcast_in_dim %6068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6070 = stablehlo.subtract %6069, %6067 : tensor<1xf32> loc(#loc5260)
    %6071 = stablehlo.divide %6064, %6070 : tensor<1xf32> loc(#loc5253)
    %6072 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6073 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6074 = stablehlo.broadcast_in_dim %6073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6075 = stablehlo.add %6072, %6074 : tensor<1xf32> loc(#loc5259)
    %6076 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6077 = stablehlo.broadcast_in_dim %6076, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6078 = stablehlo.subtract %6075, %6077 : tensor<1xf32> loc(#loc5260)
    %6079 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6080 = stablehlo.broadcast_in_dim %6079, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6081 = stablehlo.power %6080, %6078 : tensor<1xf32> loc(#loc5261)
    %6082 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6083 = stablehlo.broadcast_in_dim %6082, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6084 = stablehlo.subtract %6083, %6081 : tensor<1xf32> loc(#loc5260)
    %6085 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6086 = stablehlo.broadcast_in_dim %6085, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6087 = stablehlo.multiply %6086, %6084 : tensor<1xf32> loc(#loc5236)
    %6088 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6089 = stablehlo.broadcast_in_dim %6088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6090 = stablehlo.power %6089, %6075 : tensor<1xf32> loc(#loc5261)
    %6091 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6092 = stablehlo.broadcast_in_dim %6091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6093 = stablehlo.subtract %6092, %6090 : tensor<1xf32> loc(#loc5260)
    %6094 = stablehlo.divide %6087, %6093 : tensor<1xf32> loc(#loc5253)
    %6095 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6096 = stablehlo.broadcast_in_dim %6095, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6097 = stablehlo.subtract %6096, %6071 : tensor<1xf32> loc(#loc5260)
    %6098 = stablehlo.broadcast_in_dim %6097, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6099 = stablehlo.broadcast_in_dim %6098, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6100 = stablehlo.multiply %6099, %4340 : tensor<1x2048xf32> loc(#loc5262)
    %6101 = stablehlo.broadcast_in_dim %6071, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6102 = stablehlo.broadcast_in_dim %6101, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6103 = stablehlo.multiply %6102, %arg215 : tensor<1x2048xf32> loc(#loc5262)
    %6104 = stablehlo.add %6100, %6103 : tensor<1x2048xf32> loc(#loc5263)
    %6105 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6106 = stablehlo.broadcast_in_dim %6105, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6107 = stablehlo.subtract %6106, %6094 : tensor<1xf32> loc(#loc5260)
    %6108 = stablehlo.multiply %4340, %4340 : tensor<1x2048xf32> loc(#loc5264)
    %6109 = stablehlo.broadcast_in_dim %6107, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6110 = stablehlo.broadcast_in_dim %6109, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6111 = stablehlo.multiply %6110, %6108 : tensor<1x2048xf32> loc(#loc5262)
    %6112 = stablehlo.broadcast_in_dim %6094, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6113 = stablehlo.broadcast_in_dim %6112, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6114 = stablehlo.multiply %6113, %arg391 : tensor<1x2048xf32> loc(#loc5262)
    %6115 = stablehlo.add %6111, %6114 : tensor<1x2048xf32> loc(#loc5263)
    %6116 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6117 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6118 = stablehlo.broadcast_in_dim %6117, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6119 = stablehlo.add %6116, %6118 : tensor<1xf32> loc(#loc5259)
    %6120 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6121 = stablehlo.broadcast_in_dim %6120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6122 = stablehlo.subtract %6119, %6121 : tensor<1xf32> loc(#loc5260)
    %6123 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6124 = stablehlo.broadcast_in_dim %6123, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6125 = stablehlo.power %6124, %6122 : tensor<1xf32> loc(#loc5261)
    %6126 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6127 = stablehlo.broadcast_in_dim %6126, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6128 = stablehlo.subtract %6127, %6125 : tensor<1xf32> loc(#loc5260)
    %6129 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6130 = stablehlo.broadcast_in_dim %6129, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6131 = stablehlo.multiply %6130, %6128 : tensor<1xf32> loc(#loc5236)
    %6132 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6133 = stablehlo.broadcast_in_dim %6132, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6134 = stablehlo.power %6133, %6119 : tensor<1xf32> loc(#loc5261)
    %6135 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6136 = stablehlo.broadcast_in_dim %6135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6137 = stablehlo.subtract %6136, %6134 : tensor<1xf32> loc(#loc5260)
    %6138 = stablehlo.divide %6131, %6137 : tensor<1xf32> loc(#loc5253)
    %6139 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6141 = stablehlo.broadcast_in_dim %6140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6142 = stablehlo.add %6139, %6141 : tensor<1xf32> loc(#loc5259)
    %6143 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6144 = stablehlo.broadcast_in_dim %6143, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6145 = stablehlo.subtract %6142, %6144 : tensor<1xf32> loc(#loc5260)
    %6146 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6147 = stablehlo.broadcast_in_dim %6146, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6148 = stablehlo.power %6147, %6145 : tensor<1xf32> loc(#loc5261)
    %6149 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6150 = stablehlo.broadcast_in_dim %6149, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6151 = stablehlo.subtract %6150, %6148 : tensor<1xf32> loc(#loc5260)
    %6152 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6153 = stablehlo.broadcast_in_dim %6152, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6154 = stablehlo.multiply %6153, %6151 : tensor<1xf32> loc(#loc5236)
    %6155 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6156 = stablehlo.broadcast_in_dim %6155, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6157 = stablehlo.power %6156, %6142 : tensor<1xf32> loc(#loc5261)
    %6158 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6159 = stablehlo.broadcast_in_dim %6158, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6160 = stablehlo.subtract %6159, %6157 : tensor<1xf32> loc(#loc5260)
    %6161 = stablehlo.divide %6154, %6160 : tensor<1xf32> loc(#loc5253)
    %6162 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6163 = stablehlo.broadcast_in_dim %6162, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6164 = stablehlo.subtract %6163, %6138 : tensor<1xf32> loc(#loc5260)
    %6165 = stablehlo.broadcast_in_dim %6164, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6166 = stablehlo.broadcast_in_dim %6165, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6167 = stablehlo.multiply %6166, %4344 : tensor<1x2048xf32> loc(#loc5262)
    %6168 = stablehlo.broadcast_in_dim %6138, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6169 = stablehlo.broadcast_in_dim %6168, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6170 = stablehlo.multiply %6169, %arg216 : tensor<1x2048xf32> loc(#loc5262)
    %6171 = stablehlo.add %6167, %6170 : tensor<1x2048xf32> loc(#loc5263)
    %6172 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6173 = stablehlo.broadcast_in_dim %6172, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6174 = stablehlo.subtract %6173, %6161 : tensor<1xf32> loc(#loc5260)
    %6175 = stablehlo.multiply %4344, %4344 : tensor<1x2048xf32> loc(#loc5264)
    %6176 = stablehlo.broadcast_in_dim %6174, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6177 = stablehlo.broadcast_in_dim %6176, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6178 = stablehlo.multiply %6177, %6175 : tensor<1x2048xf32> loc(#loc5262)
    %6179 = stablehlo.broadcast_in_dim %6161, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6180 = stablehlo.broadcast_in_dim %6179, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6181 = stablehlo.multiply %6180, %arg392 : tensor<1x2048xf32> loc(#loc5262)
    %6182 = stablehlo.add %6178, %6181 : tensor<1x2048xf32> loc(#loc5263)
    %6183 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6184 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6185 = stablehlo.broadcast_in_dim %6184, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6186 = stablehlo.add %6183, %6185 : tensor<1xf32> loc(#loc5259)
    %6187 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6188 = stablehlo.broadcast_in_dim %6187, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6189 = stablehlo.subtract %6186, %6188 : tensor<1xf32> loc(#loc5260)
    %6190 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6191 = stablehlo.broadcast_in_dim %6190, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6192 = stablehlo.power %6191, %6189 : tensor<1xf32> loc(#loc5261)
    %6193 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6194 = stablehlo.broadcast_in_dim %6193, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6195 = stablehlo.subtract %6194, %6192 : tensor<1xf32> loc(#loc5260)
    %6196 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6197 = stablehlo.broadcast_in_dim %6196, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6198 = stablehlo.multiply %6197, %6195 : tensor<1xf32> loc(#loc5236)
    %6199 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6200 = stablehlo.broadcast_in_dim %6199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6201 = stablehlo.power %6200, %6186 : tensor<1xf32> loc(#loc5261)
    %6202 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6203 = stablehlo.broadcast_in_dim %6202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6204 = stablehlo.subtract %6203, %6201 : tensor<1xf32> loc(#loc5260)
    %6205 = stablehlo.divide %6198, %6204 : tensor<1xf32> loc(#loc5253)
    %6206 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6208 = stablehlo.broadcast_in_dim %6207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6209 = stablehlo.add %6206, %6208 : tensor<1xf32> loc(#loc5259)
    %6210 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6211 = stablehlo.broadcast_in_dim %6210, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6212 = stablehlo.subtract %6209, %6211 : tensor<1xf32> loc(#loc5260)
    %6213 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6214 = stablehlo.broadcast_in_dim %6213, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6215 = stablehlo.power %6214, %6212 : tensor<1xf32> loc(#loc5261)
    %6216 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6217 = stablehlo.broadcast_in_dim %6216, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6218 = stablehlo.subtract %6217, %6215 : tensor<1xf32> loc(#loc5260)
    %6219 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6220 = stablehlo.broadcast_in_dim %6219, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6221 = stablehlo.multiply %6220, %6218 : tensor<1xf32> loc(#loc5236)
    %6222 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6223 = stablehlo.broadcast_in_dim %6222, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6224 = stablehlo.power %6223, %6209 : tensor<1xf32> loc(#loc5261)
    %6225 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6226 = stablehlo.broadcast_in_dim %6225, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6227 = stablehlo.subtract %6226, %6224 : tensor<1xf32> loc(#loc5260)
    %6228 = stablehlo.divide %6221, %6227 : tensor<1xf32> loc(#loc5253)
    %6229 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6230 = stablehlo.broadcast_in_dim %6229, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6231 = stablehlo.subtract %6230, %6205 : tensor<1xf32> loc(#loc5260)
    %6232 = stablehlo.broadcast_in_dim %6231, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6233 = stablehlo.broadcast_in_dim %6232, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6234 = stablehlo.multiply %6233, %4348 : tensor<1x2048xf32> loc(#loc5262)
    %6235 = stablehlo.broadcast_in_dim %6205, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6236 = stablehlo.broadcast_in_dim %6235, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6237 = stablehlo.multiply %6236, %arg217 : tensor<1x2048xf32> loc(#loc5262)
    %6238 = stablehlo.add %6234, %6237 : tensor<1x2048xf32> loc(#loc5263)
    %6239 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6240 = stablehlo.broadcast_in_dim %6239, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6241 = stablehlo.subtract %6240, %6228 : tensor<1xf32> loc(#loc5260)
    %6242 = stablehlo.multiply %4348, %4348 : tensor<1x2048xf32> loc(#loc5264)
    %6243 = stablehlo.broadcast_in_dim %6241, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6244 = stablehlo.broadcast_in_dim %6243, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6245 = stablehlo.multiply %6244, %6242 : tensor<1x2048xf32> loc(#loc5262)
    %6246 = stablehlo.broadcast_in_dim %6228, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6247 = stablehlo.broadcast_in_dim %6246, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6248 = stablehlo.multiply %6247, %arg393 : tensor<1x2048xf32> loc(#loc5262)
    %6249 = stablehlo.add %6245, %6248 : tensor<1x2048xf32> loc(#loc5263)
    %6250 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6251 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6252 = stablehlo.broadcast_in_dim %6251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6253 = stablehlo.add %6250, %6252 : tensor<1xf32> loc(#loc5259)
    %6254 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6255 = stablehlo.broadcast_in_dim %6254, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6256 = stablehlo.subtract %6253, %6255 : tensor<1xf32> loc(#loc5260)
    %6257 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6258 = stablehlo.broadcast_in_dim %6257, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6259 = stablehlo.power %6258, %6256 : tensor<1xf32> loc(#loc5261)
    %6260 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6261 = stablehlo.broadcast_in_dim %6260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6262 = stablehlo.subtract %6261, %6259 : tensor<1xf32> loc(#loc5260)
    %6263 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6264 = stablehlo.broadcast_in_dim %6263, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6265 = stablehlo.multiply %6264, %6262 : tensor<1xf32> loc(#loc5236)
    %6266 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6267 = stablehlo.broadcast_in_dim %6266, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6268 = stablehlo.power %6267, %6253 : tensor<1xf32> loc(#loc5261)
    %6269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6270 = stablehlo.broadcast_in_dim %6269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6271 = stablehlo.subtract %6270, %6268 : tensor<1xf32> loc(#loc5260)
    %6272 = stablehlo.divide %6265, %6271 : tensor<1xf32> loc(#loc5253)
    %6273 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6274 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6275 = stablehlo.broadcast_in_dim %6274, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6276 = stablehlo.add %6273, %6275 : tensor<1xf32> loc(#loc5259)
    %6277 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6278 = stablehlo.broadcast_in_dim %6277, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6279 = stablehlo.subtract %6276, %6278 : tensor<1xf32> loc(#loc5260)
    %6280 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6281 = stablehlo.broadcast_in_dim %6280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6282 = stablehlo.power %6281, %6279 : tensor<1xf32> loc(#loc5261)
    %6283 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6284 = stablehlo.broadcast_in_dim %6283, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6285 = stablehlo.subtract %6284, %6282 : tensor<1xf32> loc(#loc5260)
    %6286 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6287 = stablehlo.broadcast_in_dim %6286, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6288 = stablehlo.multiply %6287, %6285 : tensor<1xf32> loc(#loc5236)
    %6289 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6290 = stablehlo.broadcast_in_dim %6289, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6291 = stablehlo.power %6290, %6276 : tensor<1xf32> loc(#loc5261)
    %6292 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6293 = stablehlo.broadcast_in_dim %6292, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6294 = stablehlo.subtract %6293, %6291 : tensor<1xf32> loc(#loc5260)
    %6295 = stablehlo.divide %6288, %6294 : tensor<1xf32> loc(#loc5253)
    %6296 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6297 = stablehlo.broadcast_in_dim %6296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6298 = stablehlo.subtract %6297, %6272 : tensor<1xf32> loc(#loc5260)
    %6299 = stablehlo.broadcast_in_dim %6298, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %6300 = stablehlo.broadcast_in_dim %6299, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6301 = stablehlo.multiply %6300, %4352 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6302 = stablehlo.broadcast_in_dim %6272, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %6303 = stablehlo.broadcast_in_dim %6302, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6304 = stablehlo.multiply %6303, %arg218 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6305 = stablehlo.add %6301, %6304 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %6306 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6307 = stablehlo.broadcast_in_dim %6306, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6308 = stablehlo.subtract %6307, %6295 : tensor<1xf32> loc(#loc5260)
    %6309 = stablehlo.multiply %4352, %4352 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %6310 = stablehlo.broadcast_in_dim %6308, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %6311 = stablehlo.broadcast_in_dim %6310, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6312 = stablehlo.multiply %6311, %6309 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6313 = stablehlo.broadcast_in_dim %6295, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %6314 = stablehlo.broadcast_in_dim %6313, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6315 = stablehlo.multiply %6314, %arg394 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %6316 = stablehlo.add %6312, %6315 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %6317 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6318 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6319 = stablehlo.broadcast_in_dim %6318, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6320 = stablehlo.add %6317, %6319 : tensor<1xf32> loc(#loc5259)
    %6321 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6322 = stablehlo.broadcast_in_dim %6321, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6323 = stablehlo.subtract %6320, %6322 : tensor<1xf32> loc(#loc5260)
    %6324 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6325 = stablehlo.broadcast_in_dim %6324, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6326 = stablehlo.power %6325, %6323 : tensor<1xf32> loc(#loc5261)
    %6327 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6328 = stablehlo.broadcast_in_dim %6327, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6329 = stablehlo.subtract %6328, %6326 : tensor<1xf32> loc(#loc5260)
    %6330 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6331 = stablehlo.broadcast_in_dim %6330, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6332 = stablehlo.multiply %6331, %6329 : tensor<1xf32> loc(#loc5236)
    %6333 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6334 = stablehlo.broadcast_in_dim %6333, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6335 = stablehlo.power %6334, %6320 : tensor<1xf32> loc(#loc5261)
    %6336 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6337 = stablehlo.broadcast_in_dim %6336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6338 = stablehlo.subtract %6337, %6335 : tensor<1xf32> loc(#loc5260)
    %6339 = stablehlo.divide %6332, %6338 : tensor<1xf32> loc(#loc5253)
    %6340 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6341 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6342 = stablehlo.broadcast_in_dim %6341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6343 = stablehlo.add %6340, %6342 : tensor<1xf32> loc(#loc5259)
    %6344 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6345 = stablehlo.broadcast_in_dim %6344, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6346 = stablehlo.subtract %6343, %6345 : tensor<1xf32> loc(#loc5260)
    %6347 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6348 = stablehlo.broadcast_in_dim %6347, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6349 = stablehlo.power %6348, %6346 : tensor<1xf32> loc(#loc5261)
    %6350 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6351 = stablehlo.broadcast_in_dim %6350, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6352 = stablehlo.subtract %6351, %6349 : tensor<1xf32> loc(#loc5260)
    %6353 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6354 = stablehlo.broadcast_in_dim %6353, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6355 = stablehlo.multiply %6354, %6352 : tensor<1xf32> loc(#loc5236)
    %6356 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6357 = stablehlo.broadcast_in_dim %6356, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6358 = stablehlo.power %6357, %6343 : tensor<1xf32> loc(#loc5261)
    %6359 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6360 = stablehlo.broadcast_in_dim %6359, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6361 = stablehlo.subtract %6360, %6358 : tensor<1xf32> loc(#loc5260)
    %6362 = stablehlo.divide %6355, %6361 : tensor<1xf32> loc(#loc5253)
    %6363 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6364 = stablehlo.broadcast_in_dim %6363, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6365 = stablehlo.subtract %6364, %6339 : tensor<1xf32> loc(#loc5260)
    %6366 = stablehlo.broadcast_in_dim %6365, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6367 = stablehlo.broadcast_in_dim %6366, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %6368 = stablehlo.multiply %6367, %4356 : tensor<1x64xf32> loc(#loc5262)
    %6369 = stablehlo.broadcast_in_dim %6339, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6370 = stablehlo.broadcast_in_dim %6369, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %6371 = stablehlo.multiply %6370, %arg219 : tensor<1x64xf32> loc(#loc5262)
    %6372 = stablehlo.add %6368, %6371 : tensor<1x64xf32> loc(#loc5263)
    %6373 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6374 = stablehlo.broadcast_in_dim %6373, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6375 = stablehlo.subtract %6374, %6362 : tensor<1xf32> loc(#loc5260)
    %6376 = stablehlo.multiply %4356, %4356 : tensor<1x64xf32> loc(#loc5264)
    %6377 = stablehlo.broadcast_in_dim %6375, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6378 = stablehlo.broadcast_in_dim %6377, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %6379 = stablehlo.multiply %6378, %6376 : tensor<1x64xf32> loc(#loc5262)
    %6380 = stablehlo.broadcast_in_dim %6362, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6381 = stablehlo.broadcast_in_dim %6380, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %6382 = stablehlo.multiply %6381, %arg395 : tensor<1x64xf32> loc(#loc5262)
    %6383 = stablehlo.add %6379, %6382 : tensor<1x64xf32> loc(#loc5263)
    %6384 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6385 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6386 = stablehlo.broadcast_in_dim %6385, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6387 = stablehlo.add %6384, %6386 : tensor<1xf32> loc(#loc5259)
    %6388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6389 = stablehlo.broadcast_in_dim %6388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6390 = stablehlo.subtract %6387, %6389 : tensor<1xf32> loc(#loc5260)
    %6391 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6392 = stablehlo.broadcast_in_dim %6391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6393 = stablehlo.power %6392, %6390 : tensor<1xf32> loc(#loc5261)
    %6394 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6395 = stablehlo.broadcast_in_dim %6394, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6396 = stablehlo.subtract %6395, %6393 : tensor<1xf32> loc(#loc5260)
    %6397 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6398 = stablehlo.broadcast_in_dim %6397, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6399 = stablehlo.multiply %6398, %6396 : tensor<1xf32> loc(#loc5236)
    %6400 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6401 = stablehlo.broadcast_in_dim %6400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6402 = stablehlo.power %6401, %6387 : tensor<1xf32> loc(#loc5261)
    %6403 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6404 = stablehlo.broadcast_in_dim %6403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6405 = stablehlo.subtract %6404, %6402 : tensor<1xf32> loc(#loc5260)
    %6406 = stablehlo.divide %6399, %6405 : tensor<1xf32> loc(#loc5253)
    %6407 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6409 = stablehlo.broadcast_in_dim %6408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6410 = stablehlo.add %6407, %6409 : tensor<1xf32> loc(#loc5259)
    %6411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6412 = stablehlo.broadcast_in_dim %6411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6413 = stablehlo.subtract %6410, %6412 : tensor<1xf32> loc(#loc5260)
    %6414 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6415 = stablehlo.broadcast_in_dim %6414, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6416 = stablehlo.power %6415, %6413 : tensor<1xf32> loc(#loc5261)
    %6417 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6418 = stablehlo.broadcast_in_dim %6417, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6419 = stablehlo.subtract %6418, %6416 : tensor<1xf32> loc(#loc5260)
    %6420 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6421 = stablehlo.broadcast_in_dim %6420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6422 = stablehlo.multiply %6421, %6419 : tensor<1xf32> loc(#loc5236)
    %6423 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6424 = stablehlo.broadcast_in_dim %6423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6425 = stablehlo.power %6424, %6410 : tensor<1xf32> loc(#loc5261)
    %6426 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6427 = stablehlo.broadcast_in_dim %6426, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6428 = stablehlo.subtract %6427, %6425 : tensor<1xf32> loc(#loc5260)
    %6429 = stablehlo.divide %6422, %6428 : tensor<1xf32> loc(#loc5253)
    %6430 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6431 = stablehlo.broadcast_in_dim %6430, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6432 = stablehlo.subtract %6431, %6406 : tensor<1xf32> loc(#loc5260)
    %6433 = stablehlo.broadcast_in_dim %6432, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %6434 = stablehlo.broadcast_in_dim %6433, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %6435 = stablehlo.multiply %6434, %4360 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %6436 = stablehlo.broadcast_in_dim %6406, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %6437 = stablehlo.broadcast_in_dim %6436, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %6438 = stablehlo.multiply %6437, %arg220 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %6439 = stablehlo.add %6435, %6438 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %6440 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6441 = stablehlo.broadcast_in_dim %6440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6442 = stablehlo.subtract %6441, %6429 : tensor<1xf32> loc(#loc5260)
    %6443 = stablehlo.multiply %4360, %4360 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %6444 = stablehlo.broadcast_in_dim %6442, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %6445 = stablehlo.broadcast_in_dim %6444, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %6446 = stablehlo.multiply %6445, %6443 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %6447 = stablehlo.broadcast_in_dim %6429, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %6448 = stablehlo.broadcast_in_dim %6447, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %6449 = stablehlo.multiply %6448, %arg396 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %6450 = stablehlo.add %6446, %6449 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %6451 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6452 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6453 = stablehlo.broadcast_in_dim %6452, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6454 = stablehlo.add %6451, %6453 : tensor<1xf32> loc(#loc5259)
    %6455 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6456 = stablehlo.broadcast_in_dim %6455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6457 = stablehlo.subtract %6454, %6456 : tensor<1xf32> loc(#loc5260)
    %6458 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6459 = stablehlo.broadcast_in_dim %6458, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6460 = stablehlo.power %6459, %6457 : tensor<1xf32> loc(#loc5261)
    %6461 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6462 = stablehlo.broadcast_in_dim %6461, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6463 = stablehlo.subtract %6462, %6460 : tensor<1xf32> loc(#loc5260)
    %6464 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6465 = stablehlo.broadcast_in_dim %6464, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6466 = stablehlo.multiply %6465, %6463 : tensor<1xf32> loc(#loc5236)
    %6467 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6468 = stablehlo.broadcast_in_dim %6467, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6469 = stablehlo.power %6468, %6454 : tensor<1xf32> loc(#loc5261)
    %6470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6471 = stablehlo.broadcast_in_dim %6470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6472 = stablehlo.subtract %6471, %6469 : tensor<1xf32> loc(#loc5260)
    %6473 = stablehlo.divide %6466, %6472 : tensor<1xf32> loc(#loc5253)
    %6474 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6475 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6476 = stablehlo.broadcast_in_dim %6475, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6477 = stablehlo.add %6474, %6476 : tensor<1xf32> loc(#loc5259)
    %6478 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6479 = stablehlo.broadcast_in_dim %6478, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6480 = stablehlo.subtract %6477, %6479 : tensor<1xf32> loc(#loc5260)
    %6481 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6482 = stablehlo.broadcast_in_dim %6481, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6483 = stablehlo.power %6482, %6480 : tensor<1xf32> loc(#loc5261)
    %6484 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6485 = stablehlo.broadcast_in_dim %6484, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6486 = stablehlo.subtract %6485, %6483 : tensor<1xf32> loc(#loc5260)
    %6487 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6488 = stablehlo.broadcast_in_dim %6487, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6489 = stablehlo.multiply %6488, %6486 : tensor<1xf32> loc(#loc5236)
    %6490 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6491 = stablehlo.broadcast_in_dim %6490, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6492 = stablehlo.power %6491, %6477 : tensor<1xf32> loc(#loc5261)
    %6493 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6494 = stablehlo.broadcast_in_dim %6493, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6495 = stablehlo.subtract %6494, %6492 : tensor<1xf32> loc(#loc5260)
    %6496 = stablehlo.divide %6489, %6495 : tensor<1xf32> loc(#loc5253)
    %6497 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6498 = stablehlo.broadcast_in_dim %6497, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6499 = stablehlo.subtract %6498, %6473 : tensor<1xf32> loc(#loc5260)
    %6500 = stablehlo.broadcast_in_dim %6499, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6501 = stablehlo.broadcast_in_dim %6500, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %6502 = stablehlo.multiply %6501, %4364 : tensor<1x8192xf32> loc(#loc5262)
    %6503 = stablehlo.broadcast_in_dim %6473, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6504 = stablehlo.broadcast_in_dim %6503, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %6505 = stablehlo.multiply %6504, %arg221 : tensor<1x8192xf32> loc(#loc5262)
    %6506 = stablehlo.add %6502, %6505 : tensor<1x8192xf32> loc(#loc5263)
    %6507 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6508 = stablehlo.broadcast_in_dim %6507, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6509 = stablehlo.subtract %6508, %6496 : tensor<1xf32> loc(#loc5260)
    %6510 = stablehlo.multiply %4364, %4364 : tensor<1x8192xf32> loc(#loc5264)
    %6511 = stablehlo.broadcast_in_dim %6509, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6512 = stablehlo.broadcast_in_dim %6511, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %6513 = stablehlo.multiply %6512, %6510 : tensor<1x8192xf32> loc(#loc5262)
    %6514 = stablehlo.broadcast_in_dim %6496, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6515 = stablehlo.broadcast_in_dim %6514, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %6516 = stablehlo.multiply %6515, %arg397 : tensor<1x8192xf32> loc(#loc5262)
    %6517 = stablehlo.add %6513, %6516 : tensor<1x8192xf32> loc(#loc5263)
    %6518 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6519 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6520 = stablehlo.broadcast_in_dim %6519, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6521 = stablehlo.add %6518, %6520 : tensor<1xf32> loc(#loc5259)
    %6522 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6523 = stablehlo.broadcast_in_dim %6522, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6524 = stablehlo.subtract %6521, %6523 : tensor<1xf32> loc(#loc5260)
    %6525 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6526 = stablehlo.broadcast_in_dim %6525, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6527 = stablehlo.power %6526, %6524 : tensor<1xf32> loc(#loc5261)
    %6528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6529 = stablehlo.broadcast_in_dim %6528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6530 = stablehlo.subtract %6529, %6527 : tensor<1xf32> loc(#loc5260)
    %6531 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6532 = stablehlo.broadcast_in_dim %6531, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6533 = stablehlo.multiply %6532, %6530 : tensor<1xf32> loc(#loc5236)
    %6534 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6535 = stablehlo.broadcast_in_dim %6534, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6536 = stablehlo.power %6535, %6521 : tensor<1xf32> loc(#loc5261)
    %6537 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6538 = stablehlo.broadcast_in_dim %6537, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6539 = stablehlo.subtract %6538, %6536 : tensor<1xf32> loc(#loc5260)
    %6540 = stablehlo.divide %6533, %6539 : tensor<1xf32> loc(#loc5253)
    %6541 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6542 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6543 = stablehlo.broadcast_in_dim %6542, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6544 = stablehlo.add %6541, %6543 : tensor<1xf32> loc(#loc5259)
    %6545 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6546 = stablehlo.broadcast_in_dim %6545, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6547 = stablehlo.subtract %6544, %6546 : tensor<1xf32> loc(#loc5260)
    %6548 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6549 = stablehlo.broadcast_in_dim %6548, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6550 = stablehlo.power %6549, %6547 : tensor<1xf32> loc(#loc5261)
    %6551 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6552 = stablehlo.broadcast_in_dim %6551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6553 = stablehlo.subtract %6552, %6550 : tensor<1xf32> loc(#loc5260)
    %6554 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6555 = stablehlo.broadcast_in_dim %6554, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6556 = stablehlo.multiply %6555, %6553 : tensor<1xf32> loc(#loc5236)
    %6557 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6558 = stablehlo.broadcast_in_dim %6557, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6559 = stablehlo.power %6558, %6544 : tensor<1xf32> loc(#loc5261)
    %6560 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6561 = stablehlo.broadcast_in_dim %6560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6562 = stablehlo.subtract %6561, %6559 : tensor<1xf32> loc(#loc5260)
    %6563 = stablehlo.divide %6556, %6562 : tensor<1xf32> loc(#loc5253)
    %6564 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6565 = stablehlo.broadcast_in_dim %6564, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6566 = stablehlo.subtract %6565, %6540 : tensor<1xf32> loc(#loc5260)
    %6567 = stablehlo.broadcast_in_dim %6566, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6568 = stablehlo.broadcast_in_dim %6567, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %6569 = stablehlo.multiply %6568, %4368 : tensor<1x2048x8192xf32> loc(#loc5262)
    %6570 = stablehlo.broadcast_in_dim %6540, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6571 = stablehlo.broadcast_in_dim %6570, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %6572 = stablehlo.multiply %6571, %arg222 : tensor<1x2048x8192xf32> loc(#loc5262)
    %6573 = stablehlo.add %6569, %6572 : tensor<1x2048x8192xf32> loc(#loc5263)
    %6574 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6575 = stablehlo.broadcast_in_dim %6574, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6576 = stablehlo.subtract %6575, %6563 : tensor<1xf32> loc(#loc5260)
    %6577 = stablehlo.multiply %4368, %4368 : tensor<1x2048x8192xf32> loc(#loc5264)
    %6578 = stablehlo.broadcast_in_dim %6576, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6579 = stablehlo.broadcast_in_dim %6578, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %6580 = stablehlo.multiply %6579, %6577 : tensor<1x2048x8192xf32> loc(#loc5262)
    %6581 = stablehlo.broadcast_in_dim %6563, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6582 = stablehlo.broadcast_in_dim %6581, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %6583 = stablehlo.multiply %6582, %arg398 : tensor<1x2048x8192xf32> loc(#loc5262)
    %6584 = stablehlo.add %6580, %6583 : tensor<1x2048x8192xf32> loc(#loc5263)
    %6585 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6586 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6587 = stablehlo.broadcast_in_dim %6586, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6588 = stablehlo.add %6585, %6587 : tensor<1xf32> loc(#loc5259)
    %6589 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6590 = stablehlo.broadcast_in_dim %6589, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6591 = stablehlo.subtract %6588, %6590 : tensor<1xf32> loc(#loc5260)
    %6592 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6593 = stablehlo.broadcast_in_dim %6592, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6594 = stablehlo.power %6593, %6591 : tensor<1xf32> loc(#loc5261)
    %6595 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6596 = stablehlo.broadcast_in_dim %6595, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6597 = stablehlo.subtract %6596, %6594 : tensor<1xf32> loc(#loc5260)
    %6598 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6599 = stablehlo.broadcast_in_dim %6598, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6600 = stablehlo.multiply %6599, %6597 : tensor<1xf32> loc(#loc5236)
    %6601 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6602 = stablehlo.broadcast_in_dim %6601, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6603 = stablehlo.power %6602, %6588 : tensor<1xf32> loc(#loc5261)
    %6604 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6605 = stablehlo.broadcast_in_dim %6604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6606 = stablehlo.subtract %6605, %6603 : tensor<1xf32> loc(#loc5260)
    %6607 = stablehlo.divide %6600, %6606 : tensor<1xf32> loc(#loc5253)
    %6608 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6609 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6610 = stablehlo.broadcast_in_dim %6609, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6611 = stablehlo.add %6608, %6610 : tensor<1xf32> loc(#loc5259)
    %6612 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6613 = stablehlo.broadcast_in_dim %6612, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6614 = stablehlo.subtract %6611, %6613 : tensor<1xf32> loc(#loc5260)
    %6615 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6616 = stablehlo.broadcast_in_dim %6615, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6617 = stablehlo.power %6616, %6614 : tensor<1xf32> loc(#loc5261)
    %6618 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6619 = stablehlo.broadcast_in_dim %6618, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6620 = stablehlo.subtract %6619, %6617 : tensor<1xf32> loc(#loc5260)
    %6621 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6622 = stablehlo.broadcast_in_dim %6621, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6623 = stablehlo.multiply %6622, %6620 : tensor<1xf32> loc(#loc5236)
    %6624 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6625 = stablehlo.broadcast_in_dim %6624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6626 = stablehlo.power %6625, %6611 : tensor<1xf32> loc(#loc5261)
    %6627 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6628 = stablehlo.broadcast_in_dim %6627, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6629 = stablehlo.subtract %6628, %6626 : tensor<1xf32> loc(#loc5260)
    %6630 = stablehlo.divide %6623, %6629 : tensor<1xf32> loc(#loc5253)
    %6631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6632 = stablehlo.broadcast_in_dim %6631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6633 = stablehlo.subtract %6632, %6607 : tensor<1xf32> loc(#loc5260)
    %6634 = stablehlo.broadcast_in_dim %6633, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6635 = stablehlo.broadcast_in_dim %6634, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6636 = stablehlo.multiply %6635, %4372 : tensor<1x2048xf32> loc(#loc5262)
    %6637 = stablehlo.broadcast_in_dim %6607, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6638 = stablehlo.broadcast_in_dim %6637, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6639 = stablehlo.multiply %6638, %arg223 : tensor<1x2048xf32> loc(#loc5262)
    %6640 = stablehlo.add %6636, %6639 : tensor<1x2048xf32> loc(#loc5263)
    %6641 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6642 = stablehlo.broadcast_in_dim %6641, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6643 = stablehlo.subtract %6642, %6630 : tensor<1xf32> loc(#loc5260)
    %6644 = stablehlo.multiply %4372, %4372 : tensor<1x2048xf32> loc(#loc5264)
    %6645 = stablehlo.broadcast_in_dim %6643, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6646 = stablehlo.broadcast_in_dim %6645, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6647 = stablehlo.multiply %6646, %6644 : tensor<1x2048xf32> loc(#loc5262)
    %6648 = stablehlo.broadcast_in_dim %6630, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6649 = stablehlo.broadcast_in_dim %6648, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6650 = stablehlo.multiply %6649, %arg399 : tensor<1x2048xf32> loc(#loc5262)
    %6651 = stablehlo.add %6647, %6650 : tensor<1x2048xf32> loc(#loc5263)
    %6652 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6653 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6654 = stablehlo.broadcast_in_dim %6653, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6655 = stablehlo.add %6652, %6654 : tensor<1xf32> loc(#loc5259)
    %6656 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6657 = stablehlo.broadcast_in_dim %6656, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6658 = stablehlo.subtract %6655, %6657 : tensor<1xf32> loc(#loc5260)
    %6659 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6660 = stablehlo.broadcast_in_dim %6659, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6661 = stablehlo.power %6660, %6658 : tensor<1xf32> loc(#loc5261)
    %6662 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6663 = stablehlo.broadcast_in_dim %6662, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6664 = stablehlo.subtract %6663, %6661 : tensor<1xf32> loc(#loc5260)
    %6665 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6666 = stablehlo.broadcast_in_dim %6665, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6667 = stablehlo.multiply %6666, %6664 : tensor<1xf32> loc(#loc5236)
    %6668 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6669 = stablehlo.broadcast_in_dim %6668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6670 = stablehlo.power %6669, %6655 : tensor<1xf32> loc(#loc5261)
    %6671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6672 = stablehlo.broadcast_in_dim %6671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6673 = stablehlo.subtract %6672, %6670 : tensor<1xf32> loc(#loc5260)
    %6674 = stablehlo.divide %6667, %6673 : tensor<1xf32> loc(#loc5253)
    %6675 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6676 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6677 = stablehlo.broadcast_in_dim %6676, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6678 = stablehlo.add %6675, %6677 : tensor<1xf32> loc(#loc5259)
    %6679 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6680 = stablehlo.broadcast_in_dim %6679, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6681 = stablehlo.subtract %6678, %6680 : tensor<1xf32> loc(#loc5260)
    %6682 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6683 = stablehlo.broadcast_in_dim %6682, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6684 = stablehlo.power %6683, %6681 : tensor<1xf32> loc(#loc5261)
    %6685 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6686 = stablehlo.broadcast_in_dim %6685, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6687 = stablehlo.subtract %6686, %6684 : tensor<1xf32> loc(#loc5260)
    %6688 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6689 = stablehlo.broadcast_in_dim %6688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6690 = stablehlo.multiply %6689, %6687 : tensor<1xf32> loc(#loc5236)
    %6691 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6692 = stablehlo.broadcast_in_dim %6691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6693 = stablehlo.power %6692, %6678 : tensor<1xf32> loc(#loc5261)
    %6694 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6695 = stablehlo.broadcast_in_dim %6694, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6696 = stablehlo.subtract %6695, %6693 : tensor<1xf32> loc(#loc5260)
    %6697 = stablehlo.divide %6690, %6696 : tensor<1xf32> loc(#loc5253)
    %6698 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6699 = stablehlo.broadcast_in_dim %6698, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6700 = stablehlo.subtract %6699, %6674 : tensor<1xf32> loc(#loc5260)
    %6701 = stablehlo.broadcast_in_dim %6700, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6702 = stablehlo.broadcast_in_dim %6701, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %6703 = stablehlo.multiply %6702, %4376 : tensor<1x8192x2048xf32> loc(#loc5262)
    %6704 = stablehlo.broadcast_in_dim %6674, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6705 = stablehlo.broadcast_in_dim %6704, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %6706 = stablehlo.multiply %6705, %arg224 : tensor<1x8192x2048xf32> loc(#loc5262)
    %6707 = stablehlo.add %6703, %6706 : tensor<1x8192x2048xf32> loc(#loc5263)
    %6708 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6709 = stablehlo.broadcast_in_dim %6708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6710 = stablehlo.subtract %6709, %6697 : tensor<1xf32> loc(#loc5260)
    %6711 = stablehlo.multiply %4376, %4376 : tensor<1x8192x2048xf32> loc(#loc5264)
    %6712 = stablehlo.broadcast_in_dim %6710, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6713 = stablehlo.broadcast_in_dim %6712, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %6714 = stablehlo.multiply %6713, %6711 : tensor<1x8192x2048xf32> loc(#loc5262)
    %6715 = stablehlo.broadcast_in_dim %6697, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %6716 = stablehlo.broadcast_in_dim %6715, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %6717 = stablehlo.multiply %6716, %arg400 : tensor<1x8192x2048xf32> loc(#loc5262)
    %6718 = stablehlo.add %6714, %6717 : tensor<1x8192x2048xf32> loc(#loc5263)
    %6719 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6720 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6721 = stablehlo.broadcast_in_dim %6720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6722 = stablehlo.add %6719, %6721 : tensor<1xf32> loc(#loc5259)
    %6723 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6724 = stablehlo.broadcast_in_dim %6723, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6725 = stablehlo.subtract %6722, %6724 : tensor<1xf32> loc(#loc5260)
    %6726 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6727 = stablehlo.broadcast_in_dim %6726, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6728 = stablehlo.power %6727, %6725 : tensor<1xf32> loc(#loc5261)
    %6729 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6730 = stablehlo.broadcast_in_dim %6729, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6731 = stablehlo.subtract %6730, %6728 : tensor<1xf32> loc(#loc5260)
    %6732 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6733 = stablehlo.broadcast_in_dim %6732, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6734 = stablehlo.multiply %6733, %6731 : tensor<1xf32> loc(#loc5236)
    %6735 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6736 = stablehlo.broadcast_in_dim %6735, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6737 = stablehlo.power %6736, %6722 : tensor<1xf32> loc(#loc5261)
    %6738 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6739 = stablehlo.broadcast_in_dim %6738, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6740 = stablehlo.subtract %6739, %6737 : tensor<1xf32> loc(#loc5260)
    %6741 = stablehlo.divide %6734, %6740 : tensor<1xf32> loc(#loc5253)
    %6742 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6744 = stablehlo.broadcast_in_dim %6743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6745 = stablehlo.add %6742, %6744 : tensor<1xf32> loc(#loc5259)
    %6746 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6747 = stablehlo.broadcast_in_dim %6746, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6748 = stablehlo.subtract %6745, %6747 : tensor<1xf32> loc(#loc5260)
    %6749 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6750 = stablehlo.broadcast_in_dim %6749, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6751 = stablehlo.power %6750, %6748 : tensor<1xf32> loc(#loc5261)
    %6752 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6753 = stablehlo.broadcast_in_dim %6752, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6754 = stablehlo.subtract %6753, %6751 : tensor<1xf32> loc(#loc5260)
    %6755 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6756 = stablehlo.broadcast_in_dim %6755, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6757 = stablehlo.multiply %6756, %6754 : tensor<1xf32> loc(#loc5236)
    %6758 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6759 = stablehlo.broadcast_in_dim %6758, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6760 = stablehlo.power %6759, %6745 : tensor<1xf32> loc(#loc5261)
    %6761 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6762 = stablehlo.broadcast_in_dim %6761, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6763 = stablehlo.subtract %6762, %6760 : tensor<1xf32> loc(#loc5260)
    %6764 = stablehlo.divide %6757, %6763 : tensor<1xf32> loc(#loc5253)
    %6765 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6766 = stablehlo.broadcast_in_dim %6765, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6767 = stablehlo.subtract %6766, %6741 : tensor<1xf32> loc(#loc5260)
    %6768 = stablehlo.broadcast_in_dim %6767, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6769 = stablehlo.broadcast_in_dim %6768, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6770 = stablehlo.multiply %6769, %4380 : tensor<1x2048xf32> loc(#loc5262)
    %6771 = stablehlo.broadcast_in_dim %6741, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6772 = stablehlo.broadcast_in_dim %6771, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6773 = stablehlo.multiply %6772, %arg225 : tensor<1x2048xf32> loc(#loc5262)
    %6774 = stablehlo.add %6770, %6773 : tensor<1x2048xf32> loc(#loc5263)
    %6775 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6776 = stablehlo.broadcast_in_dim %6775, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6777 = stablehlo.subtract %6776, %6764 : tensor<1xf32> loc(#loc5260)
    %6778 = stablehlo.multiply %4380, %4380 : tensor<1x2048xf32> loc(#loc5264)
    %6779 = stablehlo.broadcast_in_dim %6777, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6780 = stablehlo.broadcast_in_dim %6779, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6781 = stablehlo.multiply %6780, %6778 : tensor<1x2048xf32> loc(#loc5262)
    %6782 = stablehlo.broadcast_in_dim %6764, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6783 = stablehlo.broadcast_in_dim %6782, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6784 = stablehlo.multiply %6783, %arg401 : tensor<1x2048xf32> loc(#loc5262)
    %6785 = stablehlo.add %6781, %6784 : tensor<1x2048xf32> loc(#loc5263)
    %6786 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6787 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6788 = stablehlo.broadcast_in_dim %6787, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6789 = stablehlo.add %6786, %6788 : tensor<1xf32> loc(#loc5259)
    %6790 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6791 = stablehlo.broadcast_in_dim %6790, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6792 = stablehlo.subtract %6789, %6791 : tensor<1xf32> loc(#loc5260)
    %6793 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6794 = stablehlo.broadcast_in_dim %6793, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6795 = stablehlo.power %6794, %6792 : tensor<1xf32> loc(#loc5261)
    %6796 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6797 = stablehlo.broadcast_in_dim %6796, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6798 = stablehlo.subtract %6797, %6795 : tensor<1xf32> loc(#loc5260)
    %6799 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6800 = stablehlo.broadcast_in_dim %6799, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6801 = stablehlo.multiply %6800, %6798 : tensor<1xf32> loc(#loc5236)
    %6802 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6803 = stablehlo.broadcast_in_dim %6802, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6804 = stablehlo.power %6803, %6789 : tensor<1xf32> loc(#loc5261)
    %6805 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6806 = stablehlo.broadcast_in_dim %6805, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6807 = stablehlo.subtract %6806, %6804 : tensor<1xf32> loc(#loc5260)
    %6808 = stablehlo.divide %6801, %6807 : tensor<1xf32> loc(#loc5253)
    %6809 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6810 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6811 = stablehlo.broadcast_in_dim %6810, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6812 = stablehlo.add %6809, %6811 : tensor<1xf32> loc(#loc5259)
    %6813 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6814 = stablehlo.broadcast_in_dim %6813, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6815 = stablehlo.subtract %6812, %6814 : tensor<1xf32> loc(#loc5260)
    %6816 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6817 = stablehlo.broadcast_in_dim %6816, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6818 = stablehlo.power %6817, %6815 : tensor<1xf32> loc(#loc5261)
    %6819 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6820 = stablehlo.broadcast_in_dim %6819, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6821 = stablehlo.subtract %6820, %6818 : tensor<1xf32> loc(#loc5260)
    %6822 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6823 = stablehlo.broadcast_in_dim %6822, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6824 = stablehlo.multiply %6823, %6821 : tensor<1xf32> loc(#loc5236)
    %6825 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6826 = stablehlo.broadcast_in_dim %6825, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6827 = stablehlo.power %6826, %6812 : tensor<1xf32> loc(#loc5261)
    %6828 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6829 = stablehlo.broadcast_in_dim %6828, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6830 = stablehlo.subtract %6829, %6827 : tensor<1xf32> loc(#loc5260)
    %6831 = stablehlo.divide %6824, %6830 : tensor<1xf32> loc(#loc5253)
    %6832 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6833 = stablehlo.broadcast_in_dim %6832, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6834 = stablehlo.subtract %6833, %6808 : tensor<1xf32> loc(#loc5260)
    %6835 = stablehlo.broadcast_in_dim %6834, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6836 = stablehlo.broadcast_in_dim %6835, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6837 = stablehlo.multiply %6836, %4384 : tensor<1x2048xf32> loc(#loc5262)
    %6838 = stablehlo.broadcast_in_dim %6808, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6839 = stablehlo.broadcast_in_dim %6838, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6840 = stablehlo.multiply %6839, %arg226 : tensor<1x2048xf32> loc(#loc5262)
    %6841 = stablehlo.add %6837, %6840 : tensor<1x2048xf32> loc(#loc5263)
    %6842 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6843 = stablehlo.broadcast_in_dim %6842, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6844 = stablehlo.subtract %6843, %6831 : tensor<1xf32> loc(#loc5260)
    %6845 = stablehlo.multiply %4384, %4384 : tensor<1x2048xf32> loc(#loc5264)
    %6846 = stablehlo.broadcast_in_dim %6844, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6847 = stablehlo.broadcast_in_dim %6846, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6848 = stablehlo.multiply %6847, %6845 : tensor<1x2048xf32> loc(#loc5262)
    %6849 = stablehlo.broadcast_in_dim %6831, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6850 = stablehlo.broadcast_in_dim %6849, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6851 = stablehlo.multiply %6850, %arg402 : tensor<1x2048xf32> loc(#loc5262)
    %6852 = stablehlo.add %6848, %6851 : tensor<1x2048xf32> loc(#loc5263)
    %6853 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6854 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6855 = stablehlo.broadcast_in_dim %6854, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6856 = stablehlo.add %6853, %6855 : tensor<1xf32> loc(#loc5259)
    %6857 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6858 = stablehlo.broadcast_in_dim %6857, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6859 = stablehlo.subtract %6856, %6858 : tensor<1xf32> loc(#loc5260)
    %6860 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6861 = stablehlo.broadcast_in_dim %6860, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6862 = stablehlo.power %6861, %6859 : tensor<1xf32> loc(#loc5261)
    %6863 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6864 = stablehlo.broadcast_in_dim %6863, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6865 = stablehlo.subtract %6864, %6862 : tensor<1xf32> loc(#loc5260)
    %6866 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6867 = stablehlo.broadcast_in_dim %6866, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6868 = stablehlo.multiply %6867, %6865 : tensor<1xf32> loc(#loc5236)
    %6869 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6870 = stablehlo.broadcast_in_dim %6869, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6871 = stablehlo.power %6870, %6856 : tensor<1xf32> loc(#loc5261)
    %6872 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6873 = stablehlo.broadcast_in_dim %6872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6874 = stablehlo.subtract %6873, %6871 : tensor<1xf32> loc(#loc5260)
    %6875 = stablehlo.divide %6868, %6874 : tensor<1xf32> loc(#loc5253)
    %6876 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6877 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6878 = stablehlo.broadcast_in_dim %6877, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6879 = stablehlo.add %6876, %6878 : tensor<1xf32> loc(#loc5259)
    %6880 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6881 = stablehlo.broadcast_in_dim %6880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6882 = stablehlo.subtract %6879, %6881 : tensor<1xf32> loc(#loc5260)
    %6883 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6884 = stablehlo.broadcast_in_dim %6883, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6885 = stablehlo.power %6884, %6882 : tensor<1xf32> loc(#loc5261)
    %6886 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6887 = stablehlo.broadcast_in_dim %6886, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6888 = stablehlo.subtract %6887, %6885 : tensor<1xf32> loc(#loc5260)
    %6889 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6890 = stablehlo.broadcast_in_dim %6889, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6891 = stablehlo.multiply %6890, %6888 : tensor<1xf32> loc(#loc5236)
    %6892 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6893 = stablehlo.broadcast_in_dim %6892, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6894 = stablehlo.power %6893, %6879 : tensor<1xf32> loc(#loc5261)
    %6895 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6896 = stablehlo.broadcast_in_dim %6895, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6897 = stablehlo.subtract %6896, %6894 : tensor<1xf32> loc(#loc5260)
    %6898 = stablehlo.divide %6891, %6897 : tensor<1xf32> loc(#loc5253)
    %6899 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6900 = stablehlo.broadcast_in_dim %6899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6901 = stablehlo.subtract %6900, %6875 : tensor<1xf32> loc(#loc5260)
    %6902 = stablehlo.broadcast_in_dim %6901, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6903 = stablehlo.broadcast_in_dim %6902, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6904 = stablehlo.multiply %6903, %4388 : tensor<1x2048xf32> loc(#loc5262)
    %6905 = stablehlo.broadcast_in_dim %6875, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6906 = stablehlo.broadcast_in_dim %6905, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6907 = stablehlo.multiply %6906, %arg227 : tensor<1x2048xf32> loc(#loc5262)
    %6908 = stablehlo.add %6904, %6907 : tensor<1x2048xf32> loc(#loc5263)
    %6909 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6910 = stablehlo.broadcast_in_dim %6909, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6911 = stablehlo.subtract %6910, %6898 : tensor<1xf32> loc(#loc5260)
    %6912 = stablehlo.multiply %4388, %4388 : tensor<1x2048xf32> loc(#loc5264)
    %6913 = stablehlo.broadcast_in_dim %6911, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6914 = stablehlo.broadcast_in_dim %6913, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6915 = stablehlo.multiply %6914, %6912 : tensor<1x2048xf32> loc(#loc5262)
    %6916 = stablehlo.broadcast_in_dim %6898, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6917 = stablehlo.broadcast_in_dim %6916, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6918 = stablehlo.multiply %6917, %arg403 : tensor<1x2048xf32> loc(#loc5262)
    %6919 = stablehlo.add %6915, %6918 : tensor<1x2048xf32> loc(#loc5263)
    %6920 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6921 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6922 = stablehlo.broadcast_in_dim %6921, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6923 = stablehlo.add %6920, %6922 : tensor<1xf32> loc(#loc5259)
    %6924 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6925 = stablehlo.broadcast_in_dim %6924, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6926 = stablehlo.subtract %6923, %6925 : tensor<1xf32> loc(#loc5260)
    %6927 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6928 = stablehlo.broadcast_in_dim %6927, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6929 = stablehlo.power %6928, %6926 : tensor<1xf32> loc(#loc5261)
    %6930 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6931 = stablehlo.broadcast_in_dim %6930, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6932 = stablehlo.subtract %6931, %6929 : tensor<1xf32> loc(#loc5260)
    %6933 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6934 = stablehlo.broadcast_in_dim %6933, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6935 = stablehlo.multiply %6934, %6932 : tensor<1xf32> loc(#loc5236)
    %6936 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6937 = stablehlo.broadcast_in_dim %6936, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6938 = stablehlo.power %6937, %6923 : tensor<1xf32> loc(#loc5261)
    %6939 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6940 = stablehlo.broadcast_in_dim %6939, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6941 = stablehlo.subtract %6940, %6938 : tensor<1xf32> loc(#loc5260)
    %6942 = stablehlo.divide %6935, %6941 : tensor<1xf32> loc(#loc5253)
    %6943 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6944 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6945 = stablehlo.broadcast_in_dim %6944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6946 = stablehlo.add %6943, %6945 : tensor<1xf32> loc(#loc5259)
    %6947 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6948 = stablehlo.broadcast_in_dim %6947, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6949 = stablehlo.subtract %6946, %6948 : tensor<1xf32> loc(#loc5260)
    %6950 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6951 = stablehlo.broadcast_in_dim %6950, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6952 = stablehlo.power %6951, %6949 : tensor<1xf32> loc(#loc5261)
    %6953 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6954 = stablehlo.broadcast_in_dim %6953, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6955 = stablehlo.subtract %6954, %6952 : tensor<1xf32> loc(#loc5260)
    %6956 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6957 = stablehlo.broadcast_in_dim %6956, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %6958 = stablehlo.multiply %6957, %6955 : tensor<1xf32> loc(#loc5236)
    %6959 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %6960 = stablehlo.broadcast_in_dim %6959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6961 = stablehlo.power %6960, %6946 : tensor<1xf32> loc(#loc5261)
    %6962 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6963 = stablehlo.broadcast_in_dim %6962, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6964 = stablehlo.subtract %6963, %6961 : tensor<1xf32> loc(#loc5260)
    %6965 = stablehlo.divide %6958, %6964 : tensor<1xf32> loc(#loc5253)
    %6966 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6967 = stablehlo.broadcast_in_dim %6966, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6968 = stablehlo.subtract %6967, %6942 : tensor<1xf32> loc(#loc5260)
    %6969 = stablehlo.broadcast_in_dim %6968, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6970 = stablehlo.broadcast_in_dim %6969, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6971 = stablehlo.multiply %6970, %4392 : tensor<1x2048xf32> loc(#loc5262)
    %6972 = stablehlo.broadcast_in_dim %6942, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6973 = stablehlo.broadcast_in_dim %6972, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6974 = stablehlo.multiply %6973, %arg228 : tensor<1x2048xf32> loc(#loc5262)
    %6975 = stablehlo.add %6971, %6974 : tensor<1x2048xf32> loc(#loc5263)
    %6976 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6977 = stablehlo.broadcast_in_dim %6976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6978 = stablehlo.subtract %6977, %6965 : tensor<1xf32> loc(#loc5260)
    %6979 = stablehlo.multiply %4392, %4392 : tensor<1x2048xf32> loc(#loc5264)
    %6980 = stablehlo.broadcast_in_dim %6978, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6981 = stablehlo.broadcast_in_dim %6980, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6982 = stablehlo.multiply %6981, %6979 : tensor<1x2048xf32> loc(#loc5262)
    %6983 = stablehlo.broadcast_in_dim %6965, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %6984 = stablehlo.broadcast_in_dim %6983, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %6985 = stablehlo.multiply %6984, %arg404 : tensor<1x2048xf32> loc(#loc5262)
    %6986 = stablehlo.add %6982, %6985 : tensor<1x2048xf32> loc(#loc5263)
    %6987 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %6988 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6989 = stablehlo.broadcast_in_dim %6988, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %6990 = stablehlo.add %6987, %6989 : tensor<1xf32> loc(#loc5259)
    %6991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6992 = stablehlo.broadcast_in_dim %6991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6993 = stablehlo.subtract %6990, %6992 : tensor<1xf32> loc(#loc5260)
    %6994 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %6995 = stablehlo.broadcast_in_dim %6994, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %6996 = stablehlo.power %6995, %6993 : tensor<1xf32> loc(#loc5261)
    %6997 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %6998 = stablehlo.broadcast_in_dim %6997, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %6999 = stablehlo.subtract %6998, %6996 : tensor<1xf32> loc(#loc5260)
    %7000 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7001 = stablehlo.broadcast_in_dim %7000, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7002 = stablehlo.multiply %7001, %6999 : tensor<1xf32> loc(#loc5236)
    %7003 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7004 = stablehlo.broadcast_in_dim %7003, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7005 = stablehlo.power %7004, %6990 : tensor<1xf32> loc(#loc5261)
    %7006 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7007 = stablehlo.broadcast_in_dim %7006, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7008 = stablehlo.subtract %7007, %7005 : tensor<1xf32> loc(#loc5260)
    %7009 = stablehlo.divide %7002, %7008 : tensor<1xf32> loc(#loc5253)
    %7010 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7012 = stablehlo.broadcast_in_dim %7011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7013 = stablehlo.add %7010, %7012 : tensor<1xf32> loc(#loc5259)
    %7014 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7015 = stablehlo.broadcast_in_dim %7014, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7016 = stablehlo.subtract %7013, %7015 : tensor<1xf32> loc(#loc5260)
    %7017 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7018 = stablehlo.broadcast_in_dim %7017, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7019 = stablehlo.power %7018, %7016 : tensor<1xf32> loc(#loc5261)
    %7020 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7021 = stablehlo.broadcast_in_dim %7020, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7022 = stablehlo.subtract %7021, %7019 : tensor<1xf32> loc(#loc5260)
    %7023 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7024 = stablehlo.broadcast_in_dim %7023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7025 = stablehlo.multiply %7024, %7022 : tensor<1xf32> loc(#loc5236)
    %7026 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7027 = stablehlo.broadcast_in_dim %7026, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7028 = stablehlo.power %7027, %7013 : tensor<1xf32> loc(#loc5261)
    %7029 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7030 = stablehlo.broadcast_in_dim %7029, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7031 = stablehlo.subtract %7030, %7028 : tensor<1xf32> loc(#loc5260)
    %7032 = stablehlo.divide %7025, %7031 : tensor<1xf32> loc(#loc5253)
    %7033 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7034 = stablehlo.broadcast_in_dim %7033, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7035 = stablehlo.subtract %7034, %7009 : tensor<1xf32> loc(#loc5260)
    %7036 = stablehlo.broadcast_in_dim %7035, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7037 = stablehlo.broadcast_in_dim %7036, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7038 = stablehlo.multiply %7037, %4396 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7039 = stablehlo.broadcast_in_dim %7009, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7040 = stablehlo.broadcast_in_dim %7039, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7041 = stablehlo.multiply %7040, %arg229 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7042 = stablehlo.add %7038, %7041 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %7043 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7044 = stablehlo.broadcast_in_dim %7043, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7045 = stablehlo.subtract %7044, %7032 : tensor<1xf32> loc(#loc5260)
    %7046 = stablehlo.multiply %4396, %4396 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %7047 = stablehlo.broadcast_in_dim %7045, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7048 = stablehlo.broadcast_in_dim %7047, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7049 = stablehlo.multiply %7048, %7046 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7050 = stablehlo.broadcast_in_dim %7032, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7051 = stablehlo.broadcast_in_dim %7050, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7052 = stablehlo.multiply %7051, %arg405 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7053 = stablehlo.add %7049, %7052 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %7054 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7055 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7056 = stablehlo.broadcast_in_dim %7055, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7057 = stablehlo.add %7054, %7056 : tensor<1xf32> loc(#loc5259)
    %7058 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7059 = stablehlo.broadcast_in_dim %7058, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7060 = stablehlo.subtract %7057, %7059 : tensor<1xf32> loc(#loc5260)
    %7061 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7062 = stablehlo.broadcast_in_dim %7061, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7063 = stablehlo.power %7062, %7060 : tensor<1xf32> loc(#loc5261)
    %7064 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7065 = stablehlo.broadcast_in_dim %7064, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7066 = stablehlo.subtract %7065, %7063 : tensor<1xf32> loc(#loc5260)
    %7067 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7068 = stablehlo.broadcast_in_dim %7067, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7069 = stablehlo.multiply %7068, %7066 : tensor<1xf32> loc(#loc5236)
    %7070 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7071 = stablehlo.broadcast_in_dim %7070, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7072 = stablehlo.power %7071, %7057 : tensor<1xf32> loc(#loc5261)
    %7073 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7074 = stablehlo.broadcast_in_dim %7073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7075 = stablehlo.subtract %7074, %7072 : tensor<1xf32> loc(#loc5260)
    %7076 = stablehlo.divide %7069, %7075 : tensor<1xf32> loc(#loc5253)
    %7077 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7078 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7079 = stablehlo.broadcast_in_dim %7078, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7080 = stablehlo.add %7077, %7079 : tensor<1xf32> loc(#loc5259)
    %7081 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7082 = stablehlo.broadcast_in_dim %7081, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7083 = stablehlo.subtract %7080, %7082 : tensor<1xf32> loc(#loc5260)
    %7084 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7085 = stablehlo.broadcast_in_dim %7084, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7086 = stablehlo.power %7085, %7083 : tensor<1xf32> loc(#loc5261)
    %7087 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7088 = stablehlo.broadcast_in_dim %7087, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7089 = stablehlo.subtract %7088, %7086 : tensor<1xf32> loc(#loc5260)
    %7090 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7091 = stablehlo.broadcast_in_dim %7090, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7092 = stablehlo.multiply %7091, %7089 : tensor<1xf32> loc(#loc5236)
    %7093 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7094 = stablehlo.broadcast_in_dim %7093, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7095 = stablehlo.power %7094, %7080 : tensor<1xf32> loc(#loc5261)
    %7096 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7097 = stablehlo.broadcast_in_dim %7096, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7098 = stablehlo.subtract %7097, %7095 : tensor<1xf32> loc(#loc5260)
    %7099 = stablehlo.divide %7092, %7098 : tensor<1xf32> loc(#loc5253)
    %7100 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7101 = stablehlo.broadcast_in_dim %7100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7102 = stablehlo.subtract %7101, %7076 : tensor<1xf32> loc(#loc5260)
    %7103 = stablehlo.broadcast_in_dim %7102, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7104 = stablehlo.broadcast_in_dim %7103, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7105 = stablehlo.multiply %7104, %4400 : tensor<1x64xf32> loc(#loc5262)
    %7106 = stablehlo.broadcast_in_dim %7076, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7107 = stablehlo.broadcast_in_dim %7106, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7108 = stablehlo.multiply %7107, %arg230 : tensor<1x64xf32> loc(#loc5262)
    %7109 = stablehlo.add %7105, %7108 : tensor<1x64xf32> loc(#loc5263)
    %7110 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7111 = stablehlo.broadcast_in_dim %7110, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7112 = stablehlo.subtract %7111, %7099 : tensor<1xf32> loc(#loc5260)
    %7113 = stablehlo.multiply %4400, %4400 : tensor<1x64xf32> loc(#loc5264)
    %7114 = stablehlo.broadcast_in_dim %7112, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7115 = stablehlo.broadcast_in_dim %7114, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7116 = stablehlo.multiply %7115, %7113 : tensor<1x64xf32> loc(#loc5262)
    %7117 = stablehlo.broadcast_in_dim %7099, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7118 = stablehlo.broadcast_in_dim %7117, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7119 = stablehlo.multiply %7118, %arg406 : tensor<1x64xf32> loc(#loc5262)
    %7120 = stablehlo.add %7116, %7119 : tensor<1x64xf32> loc(#loc5263)
    %7121 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7122 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7123 = stablehlo.broadcast_in_dim %7122, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7124 = stablehlo.add %7121, %7123 : tensor<1xf32> loc(#loc5259)
    %7125 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7126 = stablehlo.broadcast_in_dim %7125, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7127 = stablehlo.subtract %7124, %7126 : tensor<1xf32> loc(#loc5260)
    %7128 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7129 = stablehlo.broadcast_in_dim %7128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7130 = stablehlo.power %7129, %7127 : tensor<1xf32> loc(#loc5261)
    %7131 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7132 = stablehlo.broadcast_in_dim %7131, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7133 = stablehlo.subtract %7132, %7130 : tensor<1xf32> loc(#loc5260)
    %7134 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7135 = stablehlo.broadcast_in_dim %7134, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7136 = stablehlo.multiply %7135, %7133 : tensor<1xf32> loc(#loc5236)
    %7137 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7138 = stablehlo.broadcast_in_dim %7137, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7139 = stablehlo.power %7138, %7124 : tensor<1xf32> loc(#loc5261)
    %7140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7141 = stablehlo.broadcast_in_dim %7140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7142 = stablehlo.subtract %7141, %7139 : tensor<1xf32> loc(#loc5260)
    %7143 = stablehlo.divide %7136, %7142 : tensor<1xf32> loc(#loc5253)
    %7144 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7145 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7146 = stablehlo.broadcast_in_dim %7145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7147 = stablehlo.add %7144, %7146 : tensor<1xf32> loc(#loc5259)
    %7148 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7149 = stablehlo.broadcast_in_dim %7148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7150 = stablehlo.subtract %7147, %7149 : tensor<1xf32> loc(#loc5260)
    %7151 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7152 = stablehlo.broadcast_in_dim %7151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7153 = stablehlo.power %7152, %7150 : tensor<1xf32> loc(#loc5261)
    %7154 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7155 = stablehlo.broadcast_in_dim %7154, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7156 = stablehlo.subtract %7155, %7153 : tensor<1xf32> loc(#loc5260)
    %7157 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7158 = stablehlo.broadcast_in_dim %7157, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7159 = stablehlo.multiply %7158, %7156 : tensor<1xf32> loc(#loc5236)
    %7160 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7161 = stablehlo.broadcast_in_dim %7160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7162 = stablehlo.power %7161, %7147 : tensor<1xf32> loc(#loc5261)
    %7163 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7164 = stablehlo.broadcast_in_dim %7163, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7165 = stablehlo.subtract %7164, %7162 : tensor<1xf32> loc(#loc5260)
    %7166 = stablehlo.divide %7159, %7165 : tensor<1xf32> loc(#loc5253)
    %7167 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7168 = stablehlo.broadcast_in_dim %7167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7169 = stablehlo.subtract %7168, %7143 : tensor<1xf32> loc(#loc5260)
    %7170 = stablehlo.broadcast_in_dim %7169, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7171 = stablehlo.broadcast_in_dim %7170, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7172 = stablehlo.multiply %7171, %4404 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7173 = stablehlo.broadcast_in_dim %7143, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7174 = stablehlo.broadcast_in_dim %7173, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7175 = stablehlo.multiply %7174, %arg231 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7176 = stablehlo.add %7172, %7175 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %7177 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7178 = stablehlo.broadcast_in_dim %7177, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7179 = stablehlo.subtract %7178, %7166 : tensor<1xf32> loc(#loc5260)
    %7180 = stablehlo.multiply %4404, %4404 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %7181 = stablehlo.broadcast_in_dim %7179, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7182 = stablehlo.broadcast_in_dim %7181, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7183 = stablehlo.multiply %7182, %7180 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7184 = stablehlo.broadcast_in_dim %7166, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7185 = stablehlo.broadcast_in_dim %7184, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7186 = stablehlo.multiply %7185, %arg407 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7187 = stablehlo.add %7183, %7186 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %7188 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7189 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7190 = stablehlo.broadcast_in_dim %7189, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7191 = stablehlo.add %7188, %7190 : tensor<1xf32> loc(#loc5259)
    %7192 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7193 = stablehlo.broadcast_in_dim %7192, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7194 = stablehlo.subtract %7191, %7193 : tensor<1xf32> loc(#loc5260)
    %7195 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7196 = stablehlo.broadcast_in_dim %7195, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7197 = stablehlo.power %7196, %7194 : tensor<1xf32> loc(#loc5261)
    %7198 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7199 = stablehlo.broadcast_in_dim %7198, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7200 = stablehlo.subtract %7199, %7197 : tensor<1xf32> loc(#loc5260)
    %7201 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7202 = stablehlo.broadcast_in_dim %7201, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7203 = stablehlo.multiply %7202, %7200 : tensor<1xf32> loc(#loc5236)
    %7204 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7205 = stablehlo.broadcast_in_dim %7204, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7206 = stablehlo.power %7205, %7191 : tensor<1xf32> loc(#loc5261)
    %7207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7208 = stablehlo.broadcast_in_dim %7207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7209 = stablehlo.subtract %7208, %7206 : tensor<1xf32> loc(#loc5260)
    %7210 = stablehlo.divide %7203, %7209 : tensor<1xf32> loc(#loc5253)
    %7211 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7212 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7213 = stablehlo.broadcast_in_dim %7212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7214 = stablehlo.add %7211, %7213 : tensor<1xf32> loc(#loc5259)
    %7215 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7216 = stablehlo.broadcast_in_dim %7215, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7217 = stablehlo.subtract %7214, %7216 : tensor<1xf32> loc(#loc5260)
    %7218 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7219 = stablehlo.broadcast_in_dim %7218, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7220 = stablehlo.power %7219, %7217 : tensor<1xf32> loc(#loc5261)
    %7221 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7222 = stablehlo.broadcast_in_dim %7221, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7223 = stablehlo.subtract %7222, %7220 : tensor<1xf32> loc(#loc5260)
    %7224 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7225 = stablehlo.broadcast_in_dim %7224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7226 = stablehlo.multiply %7225, %7223 : tensor<1xf32> loc(#loc5236)
    %7227 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7228 = stablehlo.broadcast_in_dim %7227, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7229 = stablehlo.power %7228, %7214 : tensor<1xf32> loc(#loc5261)
    %7230 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7231 = stablehlo.broadcast_in_dim %7230, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7232 = stablehlo.subtract %7231, %7229 : tensor<1xf32> loc(#loc5260)
    %7233 = stablehlo.divide %7226, %7232 : tensor<1xf32> loc(#loc5253)
    %7234 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7235 = stablehlo.broadcast_in_dim %7234, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7236 = stablehlo.subtract %7235, %7210 : tensor<1xf32> loc(#loc5260)
    %7237 = stablehlo.broadcast_in_dim %7236, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7238 = stablehlo.broadcast_in_dim %7237, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7239 = stablehlo.multiply %7238, %4408 : tensor<1x8192xf32> loc(#loc5262)
    %7240 = stablehlo.broadcast_in_dim %7210, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7241 = stablehlo.broadcast_in_dim %7240, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7242 = stablehlo.multiply %7241, %arg232 : tensor<1x8192xf32> loc(#loc5262)
    %7243 = stablehlo.add %7239, %7242 : tensor<1x8192xf32> loc(#loc5263)
    %7244 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7245 = stablehlo.broadcast_in_dim %7244, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7246 = stablehlo.subtract %7245, %7233 : tensor<1xf32> loc(#loc5260)
    %7247 = stablehlo.multiply %4408, %4408 : tensor<1x8192xf32> loc(#loc5264)
    %7248 = stablehlo.broadcast_in_dim %7246, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7249 = stablehlo.broadcast_in_dim %7248, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7250 = stablehlo.multiply %7249, %7247 : tensor<1x8192xf32> loc(#loc5262)
    %7251 = stablehlo.broadcast_in_dim %7233, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7252 = stablehlo.broadcast_in_dim %7251, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7253 = stablehlo.multiply %7252, %arg408 : tensor<1x8192xf32> loc(#loc5262)
    %7254 = stablehlo.add %7250, %7253 : tensor<1x8192xf32> loc(#loc5263)
    %7255 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7256 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7257 = stablehlo.broadcast_in_dim %7256, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7258 = stablehlo.add %7255, %7257 : tensor<1xf32> loc(#loc5259)
    %7259 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7260 = stablehlo.broadcast_in_dim %7259, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7261 = stablehlo.subtract %7258, %7260 : tensor<1xf32> loc(#loc5260)
    %7262 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7263 = stablehlo.broadcast_in_dim %7262, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7264 = stablehlo.power %7263, %7261 : tensor<1xf32> loc(#loc5261)
    %7265 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7266 = stablehlo.broadcast_in_dim %7265, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7267 = stablehlo.subtract %7266, %7264 : tensor<1xf32> loc(#loc5260)
    %7268 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7269 = stablehlo.broadcast_in_dim %7268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7270 = stablehlo.multiply %7269, %7267 : tensor<1xf32> loc(#loc5236)
    %7271 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7272 = stablehlo.broadcast_in_dim %7271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7273 = stablehlo.power %7272, %7258 : tensor<1xf32> loc(#loc5261)
    %7274 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7275 = stablehlo.broadcast_in_dim %7274, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7276 = stablehlo.subtract %7275, %7273 : tensor<1xf32> loc(#loc5260)
    %7277 = stablehlo.divide %7270, %7276 : tensor<1xf32> loc(#loc5253)
    %7278 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7279 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7280 = stablehlo.broadcast_in_dim %7279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7281 = stablehlo.add %7278, %7280 : tensor<1xf32> loc(#loc5259)
    %7282 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7283 = stablehlo.broadcast_in_dim %7282, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7284 = stablehlo.subtract %7281, %7283 : tensor<1xf32> loc(#loc5260)
    %7285 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7286 = stablehlo.broadcast_in_dim %7285, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7287 = stablehlo.power %7286, %7284 : tensor<1xf32> loc(#loc5261)
    %7288 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7289 = stablehlo.broadcast_in_dim %7288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7290 = stablehlo.subtract %7289, %7287 : tensor<1xf32> loc(#loc5260)
    %7291 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7292 = stablehlo.broadcast_in_dim %7291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7293 = stablehlo.multiply %7292, %7290 : tensor<1xf32> loc(#loc5236)
    %7294 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7295 = stablehlo.broadcast_in_dim %7294, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7296 = stablehlo.power %7295, %7281 : tensor<1xf32> loc(#loc5261)
    %7297 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7298 = stablehlo.broadcast_in_dim %7297, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7299 = stablehlo.subtract %7298, %7296 : tensor<1xf32> loc(#loc5260)
    %7300 = stablehlo.divide %7293, %7299 : tensor<1xf32> loc(#loc5253)
    %7301 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7302 = stablehlo.broadcast_in_dim %7301, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7303 = stablehlo.subtract %7302, %7277 : tensor<1xf32> loc(#loc5260)
    %7304 = stablehlo.broadcast_in_dim %7303, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7305 = stablehlo.broadcast_in_dim %7304, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %7306 = stablehlo.multiply %7305, %4412 : tensor<1x2048x8192xf32> loc(#loc5262)
    %7307 = stablehlo.broadcast_in_dim %7277, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7308 = stablehlo.broadcast_in_dim %7307, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %7309 = stablehlo.multiply %7308, %arg233 : tensor<1x2048x8192xf32> loc(#loc5262)
    %7310 = stablehlo.add %7306, %7309 : tensor<1x2048x8192xf32> loc(#loc5263)
    %7311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7312 = stablehlo.broadcast_in_dim %7311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7313 = stablehlo.subtract %7312, %7300 : tensor<1xf32> loc(#loc5260)
    %7314 = stablehlo.multiply %4412, %4412 : tensor<1x2048x8192xf32> loc(#loc5264)
    %7315 = stablehlo.broadcast_in_dim %7313, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7316 = stablehlo.broadcast_in_dim %7315, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %7317 = stablehlo.multiply %7316, %7314 : tensor<1x2048x8192xf32> loc(#loc5262)
    %7318 = stablehlo.broadcast_in_dim %7300, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7319 = stablehlo.broadcast_in_dim %7318, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %7320 = stablehlo.multiply %7319, %arg409 : tensor<1x2048x8192xf32> loc(#loc5262)
    %7321 = stablehlo.add %7317, %7320 : tensor<1x2048x8192xf32> loc(#loc5263)
    %7322 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7323 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7324 = stablehlo.broadcast_in_dim %7323, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7325 = stablehlo.add %7322, %7324 : tensor<1xf32> loc(#loc5259)
    %7326 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7327 = stablehlo.broadcast_in_dim %7326, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7328 = stablehlo.subtract %7325, %7327 : tensor<1xf32> loc(#loc5260)
    %7329 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7330 = stablehlo.broadcast_in_dim %7329, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7331 = stablehlo.power %7330, %7328 : tensor<1xf32> loc(#loc5261)
    %7332 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7333 = stablehlo.broadcast_in_dim %7332, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7334 = stablehlo.subtract %7333, %7331 : tensor<1xf32> loc(#loc5260)
    %7335 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7336 = stablehlo.broadcast_in_dim %7335, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7337 = stablehlo.multiply %7336, %7334 : tensor<1xf32> loc(#loc5236)
    %7338 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7339 = stablehlo.broadcast_in_dim %7338, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7340 = stablehlo.power %7339, %7325 : tensor<1xf32> loc(#loc5261)
    %7341 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7342 = stablehlo.broadcast_in_dim %7341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7343 = stablehlo.subtract %7342, %7340 : tensor<1xf32> loc(#loc5260)
    %7344 = stablehlo.divide %7337, %7343 : tensor<1xf32> loc(#loc5253)
    %7345 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7346 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7347 = stablehlo.broadcast_in_dim %7346, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7348 = stablehlo.add %7345, %7347 : tensor<1xf32> loc(#loc5259)
    %7349 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7350 = stablehlo.broadcast_in_dim %7349, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7351 = stablehlo.subtract %7348, %7350 : tensor<1xf32> loc(#loc5260)
    %7352 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7353 = stablehlo.broadcast_in_dim %7352, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7354 = stablehlo.power %7353, %7351 : tensor<1xf32> loc(#loc5261)
    %7355 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7356 = stablehlo.broadcast_in_dim %7355, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7357 = stablehlo.subtract %7356, %7354 : tensor<1xf32> loc(#loc5260)
    %7358 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7359 = stablehlo.broadcast_in_dim %7358, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7360 = stablehlo.multiply %7359, %7357 : tensor<1xf32> loc(#loc5236)
    %7361 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7362 = stablehlo.broadcast_in_dim %7361, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7363 = stablehlo.power %7362, %7348 : tensor<1xf32> loc(#loc5261)
    %7364 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7365 = stablehlo.broadcast_in_dim %7364, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7366 = stablehlo.subtract %7365, %7363 : tensor<1xf32> loc(#loc5260)
    %7367 = stablehlo.divide %7360, %7366 : tensor<1xf32> loc(#loc5253)
    %7368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7369 = stablehlo.broadcast_in_dim %7368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7370 = stablehlo.subtract %7369, %7344 : tensor<1xf32> loc(#loc5260)
    %7371 = stablehlo.broadcast_in_dim %7370, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7372 = stablehlo.broadcast_in_dim %7371, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7373 = stablehlo.multiply %7372, %4416 : tensor<1x2048xf32> loc(#loc5262)
    %7374 = stablehlo.broadcast_in_dim %7344, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7375 = stablehlo.broadcast_in_dim %7374, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7376 = stablehlo.multiply %7375, %arg234 : tensor<1x2048xf32> loc(#loc5262)
    %7377 = stablehlo.add %7373, %7376 : tensor<1x2048xf32> loc(#loc5263)
    %7378 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7379 = stablehlo.broadcast_in_dim %7378, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7380 = stablehlo.subtract %7379, %7367 : tensor<1xf32> loc(#loc5260)
    %7381 = stablehlo.multiply %4416, %4416 : tensor<1x2048xf32> loc(#loc5264)
    %7382 = stablehlo.broadcast_in_dim %7380, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7383 = stablehlo.broadcast_in_dim %7382, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7384 = stablehlo.multiply %7383, %7381 : tensor<1x2048xf32> loc(#loc5262)
    %7385 = stablehlo.broadcast_in_dim %7367, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7386 = stablehlo.broadcast_in_dim %7385, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7387 = stablehlo.multiply %7386, %arg410 : tensor<1x2048xf32> loc(#loc5262)
    %7388 = stablehlo.add %7384, %7387 : tensor<1x2048xf32> loc(#loc5263)
    %7389 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7390 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7391 = stablehlo.broadcast_in_dim %7390, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7392 = stablehlo.add %7389, %7391 : tensor<1xf32> loc(#loc5259)
    %7393 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7394 = stablehlo.broadcast_in_dim %7393, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7395 = stablehlo.subtract %7392, %7394 : tensor<1xf32> loc(#loc5260)
    %7396 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7397 = stablehlo.broadcast_in_dim %7396, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7398 = stablehlo.power %7397, %7395 : tensor<1xf32> loc(#loc5261)
    %7399 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7400 = stablehlo.broadcast_in_dim %7399, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7401 = stablehlo.subtract %7400, %7398 : tensor<1xf32> loc(#loc5260)
    %7402 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7403 = stablehlo.broadcast_in_dim %7402, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7404 = stablehlo.multiply %7403, %7401 : tensor<1xf32> loc(#loc5236)
    %7405 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7406 = stablehlo.broadcast_in_dim %7405, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7407 = stablehlo.power %7406, %7392 : tensor<1xf32> loc(#loc5261)
    %7408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7409 = stablehlo.broadcast_in_dim %7408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7410 = stablehlo.subtract %7409, %7407 : tensor<1xf32> loc(#loc5260)
    %7411 = stablehlo.divide %7404, %7410 : tensor<1xf32> loc(#loc5253)
    %7412 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7413 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7414 = stablehlo.broadcast_in_dim %7413, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7415 = stablehlo.add %7412, %7414 : tensor<1xf32> loc(#loc5259)
    %7416 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7417 = stablehlo.broadcast_in_dim %7416, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7418 = stablehlo.subtract %7415, %7417 : tensor<1xf32> loc(#loc5260)
    %7419 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7420 = stablehlo.broadcast_in_dim %7419, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7421 = stablehlo.power %7420, %7418 : tensor<1xf32> loc(#loc5261)
    %7422 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7423 = stablehlo.broadcast_in_dim %7422, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7424 = stablehlo.subtract %7423, %7421 : tensor<1xf32> loc(#loc5260)
    %7425 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7426 = stablehlo.broadcast_in_dim %7425, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7427 = stablehlo.multiply %7426, %7424 : tensor<1xf32> loc(#loc5236)
    %7428 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7429 = stablehlo.broadcast_in_dim %7428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7430 = stablehlo.power %7429, %7415 : tensor<1xf32> loc(#loc5261)
    %7431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7432 = stablehlo.broadcast_in_dim %7431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7433 = stablehlo.subtract %7432, %7430 : tensor<1xf32> loc(#loc5260)
    %7434 = stablehlo.divide %7427, %7433 : tensor<1xf32> loc(#loc5253)
    %7435 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7436 = stablehlo.broadcast_in_dim %7435, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7437 = stablehlo.subtract %7436, %7411 : tensor<1xf32> loc(#loc5260)
    %7438 = stablehlo.broadcast_in_dim %7437, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7439 = stablehlo.broadcast_in_dim %7438, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %7440 = stablehlo.multiply %7439, %4420 : tensor<1x8192x2048xf32> loc(#loc5262)
    %7441 = stablehlo.broadcast_in_dim %7411, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7442 = stablehlo.broadcast_in_dim %7441, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %7443 = stablehlo.multiply %7442, %arg235 : tensor<1x8192x2048xf32> loc(#loc5262)
    %7444 = stablehlo.add %7440, %7443 : tensor<1x8192x2048xf32> loc(#loc5263)
    %7445 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7446 = stablehlo.broadcast_in_dim %7445, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7447 = stablehlo.subtract %7446, %7434 : tensor<1xf32> loc(#loc5260)
    %7448 = stablehlo.multiply %4420, %4420 : tensor<1x8192x2048xf32> loc(#loc5264)
    %7449 = stablehlo.broadcast_in_dim %7447, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7450 = stablehlo.broadcast_in_dim %7449, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %7451 = stablehlo.multiply %7450, %7448 : tensor<1x8192x2048xf32> loc(#loc5262)
    %7452 = stablehlo.broadcast_in_dim %7434, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %7453 = stablehlo.broadcast_in_dim %7452, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %7454 = stablehlo.multiply %7453, %arg411 : tensor<1x8192x2048xf32> loc(#loc5262)
    %7455 = stablehlo.add %7451, %7454 : tensor<1x8192x2048xf32> loc(#loc5263)
    %7456 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7457 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7458 = stablehlo.broadcast_in_dim %7457, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7459 = stablehlo.add %7456, %7458 : tensor<1xf32> loc(#loc5259)
    %7460 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7461 = stablehlo.broadcast_in_dim %7460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7462 = stablehlo.subtract %7459, %7461 : tensor<1xf32> loc(#loc5260)
    %7463 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7464 = stablehlo.broadcast_in_dim %7463, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7465 = stablehlo.power %7464, %7462 : tensor<1xf32> loc(#loc5261)
    %7466 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7467 = stablehlo.broadcast_in_dim %7466, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7468 = stablehlo.subtract %7467, %7465 : tensor<1xf32> loc(#loc5260)
    %7469 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7470 = stablehlo.broadcast_in_dim %7469, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7471 = stablehlo.multiply %7470, %7468 : tensor<1xf32> loc(#loc5236)
    %7472 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7473 = stablehlo.broadcast_in_dim %7472, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7474 = stablehlo.power %7473, %7459 : tensor<1xf32> loc(#loc5261)
    %7475 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7476 = stablehlo.broadcast_in_dim %7475, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7477 = stablehlo.subtract %7476, %7474 : tensor<1xf32> loc(#loc5260)
    %7478 = stablehlo.divide %7471, %7477 : tensor<1xf32> loc(#loc5253)
    %7479 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7480 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7481 = stablehlo.broadcast_in_dim %7480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7482 = stablehlo.add %7479, %7481 : tensor<1xf32> loc(#loc5259)
    %7483 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7484 = stablehlo.broadcast_in_dim %7483, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7485 = stablehlo.subtract %7482, %7484 : tensor<1xf32> loc(#loc5260)
    %7486 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7487 = stablehlo.broadcast_in_dim %7486, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7488 = stablehlo.power %7487, %7485 : tensor<1xf32> loc(#loc5261)
    %7489 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7490 = stablehlo.broadcast_in_dim %7489, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7491 = stablehlo.subtract %7490, %7488 : tensor<1xf32> loc(#loc5260)
    %7492 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7493 = stablehlo.broadcast_in_dim %7492, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7494 = stablehlo.multiply %7493, %7491 : tensor<1xf32> loc(#loc5236)
    %7495 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7496 = stablehlo.broadcast_in_dim %7495, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7497 = stablehlo.power %7496, %7482 : tensor<1xf32> loc(#loc5261)
    %7498 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7499 = stablehlo.broadcast_in_dim %7498, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7500 = stablehlo.subtract %7499, %7497 : tensor<1xf32> loc(#loc5260)
    %7501 = stablehlo.divide %7494, %7500 : tensor<1xf32> loc(#loc5253)
    %7502 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7503 = stablehlo.broadcast_in_dim %7502, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7504 = stablehlo.subtract %7503, %7478 : tensor<1xf32> loc(#loc5260)
    %7505 = stablehlo.broadcast_in_dim %7504, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7506 = stablehlo.broadcast_in_dim %7505, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7507 = stablehlo.multiply %7506, %4424 : tensor<1x2048xf32> loc(#loc5262)
    %7508 = stablehlo.broadcast_in_dim %7478, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7509 = stablehlo.broadcast_in_dim %7508, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7510 = stablehlo.multiply %7509, %arg236 : tensor<1x2048xf32> loc(#loc5262)
    %7511 = stablehlo.add %7507, %7510 : tensor<1x2048xf32> loc(#loc5263)
    %7512 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7513 = stablehlo.broadcast_in_dim %7512, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7514 = stablehlo.subtract %7513, %7501 : tensor<1xf32> loc(#loc5260)
    %7515 = stablehlo.multiply %4424, %4424 : tensor<1x2048xf32> loc(#loc5264)
    %7516 = stablehlo.broadcast_in_dim %7514, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7517 = stablehlo.broadcast_in_dim %7516, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7518 = stablehlo.multiply %7517, %7515 : tensor<1x2048xf32> loc(#loc5262)
    %7519 = stablehlo.broadcast_in_dim %7501, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7520 = stablehlo.broadcast_in_dim %7519, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7521 = stablehlo.multiply %7520, %arg412 : tensor<1x2048xf32> loc(#loc5262)
    %7522 = stablehlo.add %7518, %7521 : tensor<1x2048xf32> loc(#loc5263)
    %7523 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7524 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7525 = stablehlo.broadcast_in_dim %7524, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7526 = stablehlo.add %7523, %7525 : tensor<1xf32> loc(#loc5259)
    %7527 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7528 = stablehlo.broadcast_in_dim %7527, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7529 = stablehlo.subtract %7526, %7528 : tensor<1xf32> loc(#loc5260)
    %7530 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7531 = stablehlo.broadcast_in_dim %7530, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7532 = stablehlo.power %7531, %7529 : tensor<1xf32> loc(#loc5261)
    %7533 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7534 = stablehlo.broadcast_in_dim %7533, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7535 = stablehlo.subtract %7534, %7532 : tensor<1xf32> loc(#loc5260)
    %7536 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7537 = stablehlo.broadcast_in_dim %7536, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7538 = stablehlo.multiply %7537, %7535 : tensor<1xf32> loc(#loc5236)
    %7539 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7540 = stablehlo.broadcast_in_dim %7539, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7541 = stablehlo.power %7540, %7526 : tensor<1xf32> loc(#loc5261)
    %7542 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7543 = stablehlo.broadcast_in_dim %7542, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7544 = stablehlo.subtract %7543, %7541 : tensor<1xf32> loc(#loc5260)
    %7545 = stablehlo.divide %7538, %7544 : tensor<1xf32> loc(#loc5253)
    %7546 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7547 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7548 = stablehlo.broadcast_in_dim %7547, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7549 = stablehlo.add %7546, %7548 : tensor<1xf32> loc(#loc5259)
    %7550 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7551 = stablehlo.broadcast_in_dim %7550, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7552 = stablehlo.subtract %7549, %7551 : tensor<1xf32> loc(#loc5260)
    %7553 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7554 = stablehlo.broadcast_in_dim %7553, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7555 = stablehlo.power %7554, %7552 : tensor<1xf32> loc(#loc5261)
    %7556 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7557 = stablehlo.broadcast_in_dim %7556, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7558 = stablehlo.subtract %7557, %7555 : tensor<1xf32> loc(#loc5260)
    %7559 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7560 = stablehlo.broadcast_in_dim %7559, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7561 = stablehlo.multiply %7560, %7558 : tensor<1xf32> loc(#loc5236)
    %7562 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7563 = stablehlo.broadcast_in_dim %7562, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7564 = stablehlo.power %7563, %7549 : tensor<1xf32> loc(#loc5261)
    %7565 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7566 = stablehlo.broadcast_in_dim %7565, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7567 = stablehlo.subtract %7566, %7564 : tensor<1xf32> loc(#loc5260)
    %7568 = stablehlo.divide %7561, %7567 : tensor<1xf32> loc(#loc5253)
    %7569 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7570 = stablehlo.broadcast_in_dim %7569, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7571 = stablehlo.subtract %7570, %7545 : tensor<1xf32> loc(#loc5260)
    %7572 = stablehlo.broadcast_in_dim %7571, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7573 = stablehlo.broadcast_in_dim %7572, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7574 = stablehlo.multiply %7573, %4428 : tensor<1x2048xf32> loc(#loc5262)
    %7575 = stablehlo.broadcast_in_dim %7545, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7576 = stablehlo.broadcast_in_dim %7575, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7577 = stablehlo.multiply %7576, %arg237 : tensor<1x2048xf32> loc(#loc5262)
    %7578 = stablehlo.add %7574, %7577 : tensor<1x2048xf32> loc(#loc5263)
    %7579 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7580 = stablehlo.broadcast_in_dim %7579, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7581 = stablehlo.subtract %7580, %7568 : tensor<1xf32> loc(#loc5260)
    %7582 = stablehlo.multiply %4428, %4428 : tensor<1x2048xf32> loc(#loc5264)
    %7583 = stablehlo.broadcast_in_dim %7581, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7584 = stablehlo.broadcast_in_dim %7583, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7585 = stablehlo.multiply %7584, %7582 : tensor<1x2048xf32> loc(#loc5262)
    %7586 = stablehlo.broadcast_in_dim %7568, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7587 = stablehlo.broadcast_in_dim %7586, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7588 = stablehlo.multiply %7587, %arg413 : tensor<1x2048xf32> loc(#loc5262)
    %7589 = stablehlo.add %7585, %7588 : tensor<1x2048xf32> loc(#loc5263)
    %7590 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7592 = stablehlo.broadcast_in_dim %7591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7593 = stablehlo.add %7590, %7592 : tensor<1xf32> loc(#loc5259)
    %7594 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7595 = stablehlo.broadcast_in_dim %7594, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7596 = stablehlo.subtract %7593, %7595 : tensor<1xf32> loc(#loc5260)
    %7597 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7598 = stablehlo.broadcast_in_dim %7597, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7599 = stablehlo.power %7598, %7596 : tensor<1xf32> loc(#loc5261)
    %7600 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7601 = stablehlo.broadcast_in_dim %7600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7602 = stablehlo.subtract %7601, %7599 : tensor<1xf32> loc(#loc5260)
    %7603 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7604 = stablehlo.broadcast_in_dim %7603, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7605 = stablehlo.multiply %7604, %7602 : tensor<1xf32> loc(#loc5236)
    %7606 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7607 = stablehlo.broadcast_in_dim %7606, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7608 = stablehlo.power %7607, %7593 : tensor<1xf32> loc(#loc5261)
    %7609 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7610 = stablehlo.broadcast_in_dim %7609, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7611 = stablehlo.subtract %7610, %7608 : tensor<1xf32> loc(#loc5260)
    %7612 = stablehlo.divide %7605, %7611 : tensor<1xf32> loc(#loc5253)
    %7613 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7614 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7615 = stablehlo.broadcast_in_dim %7614, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7616 = stablehlo.add %7613, %7615 : tensor<1xf32> loc(#loc5259)
    %7617 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7618 = stablehlo.broadcast_in_dim %7617, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7619 = stablehlo.subtract %7616, %7618 : tensor<1xf32> loc(#loc5260)
    %7620 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7621 = stablehlo.broadcast_in_dim %7620, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7622 = stablehlo.power %7621, %7619 : tensor<1xf32> loc(#loc5261)
    %7623 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7624 = stablehlo.broadcast_in_dim %7623, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7625 = stablehlo.subtract %7624, %7622 : tensor<1xf32> loc(#loc5260)
    %7626 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7627 = stablehlo.broadcast_in_dim %7626, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7628 = stablehlo.multiply %7627, %7625 : tensor<1xf32> loc(#loc5236)
    %7629 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7630 = stablehlo.broadcast_in_dim %7629, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7631 = stablehlo.power %7630, %7616 : tensor<1xf32> loc(#loc5261)
    %7632 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7633 = stablehlo.broadcast_in_dim %7632, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7634 = stablehlo.subtract %7633, %7631 : tensor<1xf32> loc(#loc5260)
    %7635 = stablehlo.divide %7628, %7634 : tensor<1xf32> loc(#loc5253)
    %7636 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7637 = stablehlo.broadcast_in_dim %7636, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7638 = stablehlo.subtract %7637, %7612 : tensor<1xf32> loc(#loc5260)
    %7639 = stablehlo.broadcast_in_dim %7638, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7640 = stablehlo.broadcast_in_dim %7639, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7641 = stablehlo.multiply %7640, %4432 : tensor<1x2048xf32> loc(#loc5262)
    %7642 = stablehlo.broadcast_in_dim %7612, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7643 = stablehlo.broadcast_in_dim %7642, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7644 = stablehlo.multiply %7643, %arg238 : tensor<1x2048xf32> loc(#loc5262)
    %7645 = stablehlo.add %7641, %7644 : tensor<1x2048xf32> loc(#loc5263)
    %7646 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7647 = stablehlo.broadcast_in_dim %7646, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7648 = stablehlo.subtract %7647, %7635 : tensor<1xf32> loc(#loc5260)
    %7649 = stablehlo.multiply %4432, %4432 : tensor<1x2048xf32> loc(#loc5264)
    %7650 = stablehlo.broadcast_in_dim %7648, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7651 = stablehlo.broadcast_in_dim %7650, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7652 = stablehlo.multiply %7651, %7649 : tensor<1x2048xf32> loc(#loc5262)
    %7653 = stablehlo.broadcast_in_dim %7635, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7654 = stablehlo.broadcast_in_dim %7653, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7655 = stablehlo.multiply %7654, %arg414 : tensor<1x2048xf32> loc(#loc5262)
    %7656 = stablehlo.add %7652, %7655 : tensor<1x2048xf32> loc(#loc5263)
    %7657 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7658 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7659 = stablehlo.broadcast_in_dim %7658, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7660 = stablehlo.add %7657, %7659 : tensor<1xf32> loc(#loc5259)
    %7661 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7662 = stablehlo.broadcast_in_dim %7661, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7663 = stablehlo.subtract %7660, %7662 : tensor<1xf32> loc(#loc5260)
    %7664 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7665 = stablehlo.broadcast_in_dim %7664, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7666 = stablehlo.power %7665, %7663 : tensor<1xf32> loc(#loc5261)
    %7667 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7668 = stablehlo.broadcast_in_dim %7667, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7669 = stablehlo.subtract %7668, %7666 : tensor<1xf32> loc(#loc5260)
    %7670 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7671 = stablehlo.broadcast_in_dim %7670, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7672 = stablehlo.multiply %7671, %7669 : tensor<1xf32> loc(#loc5236)
    %7673 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7674 = stablehlo.broadcast_in_dim %7673, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7675 = stablehlo.power %7674, %7660 : tensor<1xf32> loc(#loc5261)
    %7676 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7677 = stablehlo.broadcast_in_dim %7676, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7678 = stablehlo.subtract %7677, %7675 : tensor<1xf32> loc(#loc5260)
    %7679 = stablehlo.divide %7672, %7678 : tensor<1xf32> loc(#loc5253)
    %7680 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7681 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7682 = stablehlo.broadcast_in_dim %7681, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7683 = stablehlo.add %7680, %7682 : tensor<1xf32> loc(#loc5259)
    %7684 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7685 = stablehlo.broadcast_in_dim %7684, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7686 = stablehlo.subtract %7683, %7685 : tensor<1xf32> loc(#loc5260)
    %7687 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7688 = stablehlo.broadcast_in_dim %7687, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7689 = stablehlo.power %7688, %7686 : tensor<1xf32> loc(#loc5261)
    %7690 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7691 = stablehlo.broadcast_in_dim %7690, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7692 = stablehlo.subtract %7691, %7689 : tensor<1xf32> loc(#loc5260)
    %7693 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7694 = stablehlo.broadcast_in_dim %7693, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7695 = stablehlo.multiply %7694, %7692 : tensor<1xf32> loc(#loc5236)
    %7696 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7697 = stablehlo.broadcast_in_dim %7696, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7698 = stablehlo.power %7697, %7683 : tensor<1xf32> loc(#loc5261)
    %7699 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7700 = stablehlo.broadcast_in_dim %7699, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7701 = stablehlo.subtract %7700, %7698 : tensor<1xf32> loc(#loc5260)
    %7702 = stablehlo.divide %7695, %7701 : tensor<1xf32> loc(#loc5253)
    %7703 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7704 = stablehlo.broadcast_in_dim %7703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7705 = stablehlo.subtract %7704, %7679 : tensor<1xf32> loc(#loc5260)
    %7706 = stablehlo.broadcast_in_dim %7705, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7707 = stablehlo.broadcast_in_dim %7706, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7708 = stablehlo.multiply %7707, %4436 : tensor<1x2048xf32> loc(#loc5262)
    %7709 = stablehlo.broadcast_in_dim %7679, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7710 = stablehlo.broadcast_in_dim %7709, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7711 = stablehlo.multiply %7710, %arg239 : tensor<1x2048xf32> loc(#loc5262)
    %7712 = stablehlo.add %7708, %7711 : tensor<1x2048xf32> loc(#loc5263)
    %7713 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7714 = stablehlo.broadcast_in_dim %7713, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7715 = stablehlo.subtract %7714, %7702 : tensor<1xf32> loc(#loc5260)
    %7716 = stablehlo.multiply %4436, %4436 : tensor<1x2048xf32> loc(#loc5264)
    %7717 = stablehlo.broadcast_in_dim %7715, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7718 = stablehlo.broadcast_in_dim %7717, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7719 = stablehlo.multiply %7718, %7716 : tensor<1x2048xf32> loc(#loc5262)
    %7720 = stablehlo.broadcast_in_dim %7702, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7721 = stablehlo.broadcast_in_dim %7720, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %7722 = stablehlo.multiply %7721, %arg415 : tensor<1x2048xf32> loc(#loc5262)
    %7723 = stablehlo.add %7719, %7722 : tensor<1x2048xf32> loc(#loc5263)
    %7724 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7725 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7726 = stablehlo.broadcast_in_dim %7725, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7727 = stablehlo.add %7724, %7726 : tensor<1xf32> loc(#loc5259)
    %7728 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7729 = stablehlo.broadcast_in_dim %7728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7730 = stablehlo.subtract %7727, %7729 : tensor<1xf32> loc(#loc5260)
    %7731 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7732 = stablehlo.broadcast_in_dim %7731, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7733 = stablehlo.power %7732, %7730 : tensor<1xf32> loc(#loc5261)
    %7734 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7735 = stablehlo.broadcast_in_dim %7734, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7736 = stablehlo.subtract %7735, %7733 : tensor<1xf32> loc(#loc5260)
    %7737 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7738 = stablehlo.broadcast_in_dim %7737, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7739 = stablehlo.multiply %7738, %7736 : tensor<1xf32> loc(#loc5236)
    %7740 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7741 = stablehlo.broadcast_in_dim %7740, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7742 = stablehlo.power %7741, %7727 : tensor<1xf32> loc(#loc5261)
    %7743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7744 = stablehlo.broadcast_in_dim %7743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7745 = stablehlo.subtract %7744, %7742 : tensor<1xf32> loc(#loc5260)
    %7746 = stablehlo.divide %7739, %7745 : tensor<1xf32> loc(#loc5253)
    %7747 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7749 = stablehlo.broadcast_in_dim %7748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7750 = stablehlo.add %7747, %7749 : tensor<1xf32> loc(#loc5259)
    %7751 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7752 = stablehlo.broadcast_in_dim %7751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7753 = stablehlo.subtract %7750, %7752 : tensor<1xf32> loc(#loc5260)
    %7754 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7755 = stablehlo.broadcast_in_dim %7754, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7756 = stablehlo.power %7755, %7753 : tensor<1xf32> loc(#loc5261)
    %7757 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7758 = stablehlo.broadcast_in_dim %7757, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7759 = stablehlo.subtract %7758, %7756 : tensor<1xf32> loc(#loc5260)
    %7760 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7761 = stablehlo.broadcast_in_dim %7760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7762 = stablehlo.multiply %7761, %7759 : tensor<1xf32> loc(#loc5236)
    %7763 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7764 = stablehlo.broadcast_in_dim %7763, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7765 = stablehlo.power %7764, %7750 : tensor<1xf32> loc(#loc5261)
    %7766 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7767 = stablehlo.broadcast_in_dim %7766, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7768 = stablehlo.subtract %7767, %7765 : tensor<1xf32> loc(#loc5260)
    %7769 = stablehlo.divide %7762, %7768 : tensor<1xf32> loc(#loc5253)
    %7770 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7771 = stablehlo.broadcast_in_dim %7770, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7772 = stablehlo.subtract %7771, %7746 : tensor<1xf32> loc(#loc5260)
    %7773 = stablehlo.broadcast_in_dim %7772, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7774 = stablehlo.broadcast_in_dim %7773, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7775 = stablehlo.multiply %7774, %4440 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7776 = stablehlo.broadcast_in_dim %7746, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7777 = stablehlo.broadcast_in_dim %7776, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7778 = stablehlo.multiply %7777, %arg240 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7779 = stablehlo.add %7775, %7778 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %7780 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7781 = stablehlo.broadcast_in_dim %7780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7782 = stablehlo.subtract %7781, %7769 : tensor<1xf32> loc(#loc5260)
    %7783 = stablehlo.multiply %4440, %4440 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %7784 = stablehlo.broadcast_in_dim %7782, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7785 = stablehlo.broadcast_in_dim %7784, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7786 = stablehlo.multiply %7785, %7783 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7787 = stablehlo.broadcast_in_dim %7769, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %7788 = stablehlo.broadcast_in_dim %7787, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7789 = stablehlo.multiply %7788, %arg416 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %7790 = stablehlo.add %7786, %7789 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %7791 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7792 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7793 = stablehlo.broadcast_in_dim %7792, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7794 = stablehlo.add %7791, %7793 : tensor<1xf32> loc(#loc5259)
    %7795 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7796 = stablehlo.broadcast_in_dim %7795, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7797 = stablehlo.subtract %7794, %7796 : tensor<1xf32> loc(#loc5260)
    %7798 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7799 = stablehlo.broadcast_in_dim %7798, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7800 = stablehlo.power %7799, %7797 : tensor<1xf32> loc(#loc5261)
    %7801 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7802 = stablehlo.broadcast_in_dim %7801, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7803 = stablehlo.subtract %7802, %7800 : tensor<1xf32> loc(#loc5260)
    %7804 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7805 = stablehlo.broadcast_in_dim %7804, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7806 = stablehlo.multiply %7805, %7803 : tensor<1xf32> loc(#loc5236)
    %7807 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7808 = stablehlo.broadcast_in_dim %7807, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7809 = stablehlo.power %7808, %7794 : tensor<1xf32> loc(#loc5261)
    %7810 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7811 = stablehlo.broadcast_in_dim %7810, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7812 = stablehlo.subtract %7811, %7809 : tensor<1xf32> loc(#loc5260)
    %7813 = stablehlo.divide %7806, %7812 : tensor<1xf32> loc(#loc5253)
    %7814 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7815 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7816 = stablehlo.broadcast_in_dim %7815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7817 = stablehlo.add %7814, %7816 : tensor<1xf32> loc(#loc5259)
    %7818 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7819 = stablehlo.broadcast_in_dim %7818, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7820 = stablehlo.subtract %7817, %7819 : tensor<1xf32> loc(#loc5260)
    %7821 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7822 = stablehlo.broadcast_in_dim %7821, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7823 = stablehlo.power %7822, %7820 : tensor<1xf32> loc(#loc5261)
    %7824 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7825 = stablehlo.broadcast_in_dim %7824, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7826 = stablehlo.subtract %7825, %7823 : tensor<1xf32> loc(#loc5260)
    %7827 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7828 = stablehlo.broadcast_in_dim %7827, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7829 = stablehlo.multiply %7828, %7826 : tensor<1xf32> loc(#loc5236)
    %7830 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7831 = stablehlo.broadcast_in_dim %7830, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7832 = stablehlo.power %7831, %7817 : tensor<1xf32> loc(#loc5261)
    %7833 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7834 = stablehlo.broadcast_in_dim %7833, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7835 = stablehlo.subtract %7834, %7832 : tensor<1xf32> loc(#loc5260)
    %7836 = stablehlo.divide %7829, %7835 : tensor<1xf32> loc(#loc5253)
    %7837 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7838 = stablehlo.broadcast_in_dim %7837, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7839 = stablehlo.subtract %7838, %7813 : tensor<1xf32> loc(#loc5260)
    %7840 = stablehlo.broadcast_in_dim %7839, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7841 = stablehlo.broadcast_in_dim %7840, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7842 = stablehlo.multiply %7841, %4444 : tensor<1x64xf32> loc(#loc5262)
    %7843 = stablehlo.broadcast_in_dim %7813, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7844 = stablehlo.broadcast_in_dim %7843, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7845 = stablehlo.multiply %7844, %arg241 : tensor<1x64xf32> loc(#loc5262)
    %7846 = stablehlo.add %7842, %7845 : tensor<1x64xf32> loc(#loc5263)
    %7847 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7848 = stablehlo.broadcast_in_dim %7847, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7849 = stablehlo.subtract %7848, %7836 : tensor<1xf32> loc(#loc5260)
    %7850 = stablehlo.multiply %4444, %4444 : tensor<1x64xf32> loc(#loc5264)
    %7851 = stablehlo.broadcast_in_dim %7849, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7852 = stablehlo.broadcast_in_dim %7851, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7853 = stablehlo.multiply %7852, %7850 : tensor<1x64xf32> loc(#loc5262)
    %7854 = stablehlo.broadcast_in_dim %7836, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7855 = stablehlo.broadcast_in_dim %7854, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %7856 = stablehlo.multiply %7855, %arg417 : tensor<1x64xf32> loc(#loc5262)
    %7857 = stablehlo.add %7853, %7856 : tensor<1x64xf32> loc(#loc5263)
    %7858 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7859 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7860 = stablehlo.broadcast_in_dim %7859, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7861 = stablehlo.add %7858, %7860 : tensor<1xf32> loc(#loc5259)
    %7862 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7863 = stablehlo.broadcast_in_dim %7862, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7864 = stablehlo.subtract %7861, %7863 : tensor<1xf32> loc(#loc5260)
    %7865 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7866 = stablehlo.broadcast_in_dim %7865, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7867 = stablehlo.power %7866, %7864 : tensor<1xf32> loc(#loc5261)
    %7868 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7869 = stablehlo.broadcast_in_dim %7868, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7870 = stablehlo.subtract %7869, %7867 : tensor<1xf32> loc(#loc5260)
    %7871 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7872 = stablehlo.broadcast_in_dim %7871, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7873 = stablehlo.multiply %7872, %7870 : tensor<1xf32> loc(#loc5236)
    %7874 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7875 = stablehlo.broadcast_in_dim %7874, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7876 = stablehlo.power %7875, %7861 : tensor<1xf32> loc(#loc5261)
    %7877 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7878 = stablehlo.broadcast_in_dim %7877, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7879 = stablehlo.subtract %7878, %7876 : tensor<1xf32> loc(#loc5260)
    %7880 = stablehlo.divide %7873, %7879 : tensor<1xf32> loc(#loc5253)
    %7881 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7882 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7883 = stablehlo.broadcast_in_dim %7882, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7884 = stablehlo.add %7881, %7883 : tensor<1xf32> loc(#loc5259)
    %7885 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7886 = stablehlo.broadcast_in_dim %7885, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7887 = stablehlo.subtract %7884, %7886 : tensor<1xf32> loc(#loc5260)
    %7888 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7889 = stablehlo.broadcast_in_dim %7888, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7890 = stablehlo.power %7889, %7887 : tensor<1xf32> loc(#loc5261)
    %7891 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7892 = stablehlo.broadcast_in_dim %7891, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7893 = stablehlo.subtract %7892, %7890 : tensor<1xf32> loc(#loc5260)
    %7894 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7895 = stablehlo.broadcast_in_dim %7894, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7896 = stablehlo.multiply %7895, %7893 : tensor<1xf32> loc(#loc5236)
    %7897 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7898 = stablehlo.broadcast_in_dim %7897, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7899 = stablehlo.power %7898, %7884 : tensor<1xf32> loc(#loc5261)
    %7900 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7901 = stablehlo.broadcast_in_dim %7900, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7902 = stablehlo.subtract %7901, %7899 : tensor<1xf32> loc(#loc5260)
    %7903 = stablehlo.divide %7896, %7902 : tensor<1xf32> loc(#loc5253)
    %7904 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7905 = stablehlo.broadcast_in_dim %7904, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7906 = stablehlo.subtract %7905, %7880 : tensor<1xf32> loc(#loc5260)
    %7907 = stablehlo.broadcast_in_dim %7906, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7908 = stablehlo.broadcast_in_dim %7907, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7909 = stablehlo.multiply %7908, %4448 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7910 = stablehlo.broadcast_in_dim %7880, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7911 = stablehlo.broadcast_in_dim %7910, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7912 = stablehlo.multiply %7911, %arg242 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7913 = stablehlo.add %7909, %7912 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %7914 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7915 = stablehlo.broadcast_in_dim %7914, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7916 = stablehlo.subtract %7915, %7903 : tensor<1xf32> loc(#loc5260)
    %7917 = stablehlo.multiply %4448, %4448 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %7918 = stablehlo.broadcast_in_dim %7916, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7919 = stablehlo.broadcast_in_dim %7918, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7920 = stablehlo.multiply %7919, %7917 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7921 = stablehlo.broadcast_in_dim %7903, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %7922 = stablehlo.broadcast_in_dim %7921, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %7923 = stablehlo.multiply %7922, %arg418 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %7924 = stablehlo.add %7920, %7923 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %7925 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7926 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7927 = stablehlo.broadcast_in_dim %7926, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7928 = stablehlo.add %7925, %7927 : tensor<1xf32> loc(#loc5259)
    %7929 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7930 = stablehlo.broadcast_in_dim %7929, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7931 = stablehlo.subtract %7928, %7930 : tensor<1xf32> loc(#loc5260)
    %7932 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7933 = stablehlo.broadcast_in_dim %7932, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7934 = stablehlo.power %7933, %7931 : tensor<1xf32> loc(#loc5261)
    %7935 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7936 = stablehlo.broadcast_in_dim %7935, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7937 = stablehlo.subtract %7936, %7934 : tensor<1xf32> loc(#loc5260)
    %7938 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7939 = stablehlo.broadcast_in_dim %7938, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7940 = stablehlo.multiply %7939, %7937 : tensor<1xf32> loc(#loc5236)
    %7941 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %7942 = stablehlo.broadcast_in_dim %7941, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7943 = stablehlo.power %7942, %7928 : tensor<1xf32> loc(#loc5261)
    %7944 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7945 = stablehlo.broadcast_in_dim %7944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7946 = stablehlo.subtract %7945, %7943 : tensor<1xf32> loc(#loc5260)
    %7947 = stablehlo.divide %7940, %7946 : tensor<1xf32> loc(#loc5253)
    %7948 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7949 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7950 = stablehlo.broadcast_in_dim %7949, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7951 = stablehlo.add %7948, %7950 : tensor<1xf32> loc(#loc5259)
    %7952 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7953 = stablehlo.broadcast_in_dim %7952, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7954 = stablehlo.subtract %7951, %7953 : tensor<1xf32> loc(#loc5260)
    %7955 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7956 = stablehlo.broadcast_in_dim %7955, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7957 = stablehlo.power %7956, %7954 : tensor<1xf32> loc(#loc5261)
    %7958 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7959 = stablehlo.broadcast_in_dim %7958, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7960 = stablehlo.subtract %7959, %7957 : tensor<1xf32> loc(#loc5260)
    %7961 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7962 = stablehlo.broadcast_in_dim %7961, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %7963 = stablehlo.multiply %7962, %7960 : tensor<1xf32> loc(#loc5236)
    %7964 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %7965 = stablehlo.broadcast_in_dim %7964, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %7966 = stablehlo.power %7965, %7951 : tensor<1xf32> loc(#loc5261)
    %7967 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7968 = stablehlo.broadcast_in_dim %7967, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7969 = stablehlo.subtract %7968, %7966 : tensor<1xf32> loc(#loc5260)
    %7970 = stablehlo.divide %7963, %7969 : tensor<1xf32> loc(#loc5253)
    %7971 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7972 = stablehlo.broadcast_in_dim %7971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7973 = stablehlo.subtract %7972, %7947 : tensor<1xf32> loc(#loc5260)
    %7974 = stablehlo.broadcast_in_dim %7973, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7975 = stablehlo.broadcast_in_dim %7974, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7976 = stablehlo.multiply %7975, %4452 : tensor<1x8192xf32> loc(#loc5262)
    %7977 = stablehlo.broadcast_in_dim %7947, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7978 = stablehlo.broadcast_in_dim %7977, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7979 = stablehlo.multiply %7978, %arg243 : tensor<1x8192xf32> loc(#loc5262)
    %7980 = stablehlo.add %7976, %7979 : tensor<1x8192xf32> loc(#loc5263)
    %7981 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7982 = stablehlo.broadcast_in_dim %7981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7983 = stablehlo.subtract %7982, %7970 : tensor<1xf32> loc(#loc5260)
    %7984 = stablehlo.multiply %4452, %4452 : tensor<1x8192xf32> loc(#loc5264)
    %7985 = stablehlo.broadcast_in_dim %7983, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7986 = stablehlo.broadcast_in_dim %7985, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7987 = stablehlo.multiply %7986, %7984 : tensor<1x8192xf32> loc(#loc5262)
    %7988 = stablehlo.broadcast_in_dim %7970, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %7989 = stablehlo.broadcast_in_dim %7988, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %7990 = stablehlo.multiply %7989, %arg419 : tensor<1x8192xf32> loc(#loc5262)
    %7991 = stablehlo.add %7987, %7990 : tensor<1x8192xf32> loc(#loc5263)
    %7992 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %7993 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7994 = stablehlo.broadcast_in_dim %7993, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %7995 = stablehlo.add %7992, %7994 : tensor<1xf32> loc(#loc5259)
    %7996 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %7997 = stablehlo.broadcast_in_dim %7996, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %7998 = stablehlo.subtract %7995, %7997 : tensor<1xf32> loc(#loc5260)
    %7999 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8000 = stablehlo.broadcast_in_dim %7999, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8001 = stablehlo.power %8000, %7998 : tensor<1xf32> loc(#loc5261)
    %8002 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8003 = stablehlo.broadcast_in_dim %8002, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8004 = stablehlo.subtract %8003, %8001 : tensor<1xf32> loc(#loc5260)
    %8005 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8006 = stablehlo.broadcast_in_dim %8005, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8007 = stablehlo.multiply %8006, %8004 : tensor<1xf32> loc(#loc5236)
    %8008 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8009 = stablehlo.broadcast_in_dim %8008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8010 = stablehlo.power %8009, %7995 : tensor<1xf32> loc(#loc5261)
    %8011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8012 = stablehlo.broadcast_in_dim %8011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8013 = stablehlo.subtract %8012, %8010 : tensor<1xf32> loc(#loc5260)
    %8014 = stablehlo.divide %8007, %8013 : tensor<1xf32> loc(#loc5253)
    %8015 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8016 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8017 = stablehlo.broadcast_in_dim %8016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8018 = stablehlo.add %8015, %8017 : tensor<1xf32> loc(#loc5259)
    %8019 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8020 = stablehlo.broadcast_in_dim %8019, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8021 = stablehlo.subtract %8018, %8020 : tensor<1xf32> loc(#loc5260)
    %8022 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8023 = stablehlo.broadcast_in_dim %8022, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8024 = stablehlo.power %8023, %8021 : tensor<1xf32> loc(#loc5261)
    %8025 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8026 = stablehlo.broadcast_in_dim %8025, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8027 = stablehlo.subtract %8026, %8024 : tensor<1xf32> loc(#loc5260)
    %8028 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8029 = stablehlo.broadcast_in_dim %8028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8030 = stablehlo.multiply %8029, %8027 : tensor<1xf32> loc(#loc5236)
    %8031 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8032 = stablehlo.broadcast_in_dim %8031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8033 = stablehlo.power %8032, %8018 : tensor<1xf32> loc(#loc5261)
    %8034 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8035 = stablehlo.broadcast_in_dim %8034, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8036 = stablehlo.subtract %8035, %8033 : tensor<1xf32> loc(#loc5260)
    %8037 = stablehlo.divide %8030, %8036 : tensor<1xf32> loc(#loc5253)
    %8038 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8039 = stablehlo.broadcast_in_dim %8038, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8040 = stablehlo.subtract %8039, %8014 : tensor<1xf32> loc(#loc5260)
    %8041 = stablehlo.broadcast_in_dim %8040, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8042 = stablehlo.broadcast_in_dim %8041, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8043 = stablehlo.multiply %8042, %4456 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8044 = stablehlo.broadcast_in_dim %8014, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8045 = stablehlo.broadcast_in_dim %8044, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8046 = stablehlo.multiply %8045, %arg244 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8047 = stablehlo.add %8043, %8046 : tensor<1x2048x8192xf32> loc(#loc5263)
    %8048 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8049 = stablehlo.broadcast_in_dim %8048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8050 = stablehlo.subtract %8049, %8037 : tensor<1xf32> loc(#loc5260)
    %8051 = stablehlo.multiply %4456, %4456 : tensor<1x2048x8192xf32> loc(#loc5264)
    %8052 = stablehlo.broadcast_in_dim %8050, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8053 = stablehlo.broadcast_in_dim %8052, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8054 = stablehlo.multiply %8053, %8051 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8055 = stablehlo.broadcast_in_dim %8037, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8056 = stablehlo.broadcast_in_dim %8055, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8057 = stablehlo.multiply %8056, %arg420 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8058 = stablehlo.add %8054, %8057 : tensor<1x2048x8192xf32> loc(#loc5263)
    %8059 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8060 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8061 = stablehlo.broadcast_in_dim %8060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8062 = stablehlo.add %8059, %8061 : tensor<1xf32> loc(#loc5259)
    %8063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8064 = stablehlo.broadcast_in_dim %8063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8065 = stablehlo.subtract %8062, %8064 : tensor<1xf32> loc(#loc5260)
    %8066 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8067 = stablehlo.broadcast_in_dim %8066, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8068 = stablehlo.power %8067, %8065 : tensor<1xf32> loc(#loc5261)
    %8069 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8070 = stablehlo.broadcast_in_dim %8069, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8071 = stablehlo.subtract %8070, %8068 : tensor<1xf32> loc(#loc5260)
    %8072 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8073 = stablehlo.broadcast_in_dim %8072, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8074 = stablehlo.multiply %8073, %8071 : tensor<1xf32> loc(#loc5236)
    %8075 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8076 = stablehlo.broadcast_in_dim %8075, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8077 = stablehlo.power %8076, %8062 : tensor<1xf32> loc(#loc5261)
    %8078 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8079 = stablehlo.broadcast_in_dim %8078, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8080 = stablehlo.subtract %8079, %8077 : tensor<1xf32> loc(#loc5260)
    %8081 = stablehlo.divide %8074, %8080 : tensor<1xf32> loc(#loc5253)
    %8082 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8083 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8084 = stablehlo.broadcast_in_dim %8083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8085 = stablehlo.add %8082, %8084 : tensor<1xf32> loc(#loc5259)
    %8086 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8087 = stablehlo.broadcast_in_dim %8086, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8088 = stablehlo.subtract %8085, %8087 : tensor<1xf32> loc(#loc5260)
    %8089 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8090 = stablehlo.broadcast_in_dim %8089, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8091 = stablehlo.power %8090, %8088 : tensor<1xf32> loc(#loc5261)
    %8092 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8093 = stablehlo.broadcast_in_dim %8092, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8094 = stablehlo.subtract %8093, %8091 : tensor<1xf32> loc(#loc5260)
    %8095 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8096 = stablehlo.broadcast_in_dim %8095, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8097 = stablehlo.multiply %8096, %8094 : tensor<1xf32> loc(#loc5236)
    %8098 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8099 = stablehlo.broadcast_in_dim %8098, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8100 = stablehlo.power %8099, %8085 : tensor<1xf32> loc(#loc5261)
    %8101 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8102 = stablehlo.broadcast_in_dim %8101, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8103 = stablehlo.subtract %8102, %8100 : tensor<1xf32> loc(#loc5260)
    %8104 = stablehlo.divide %8097, %8103 : tensor<1xf32> loc(#loc5253)
    %8105 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8106 = stablehlo.broadcast_in_dim %8105, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8107 = stablehlo.subtract %8106, %8081 : tensor<1xf32> loc(#loc5260)
    %8108 = stablehlo.broadcast_in_dim %8107, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8109 = stablehlo.broadcast_in_dim %8108, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8110 = stablehlo.multiply %8109, %4460 : tensor<1x2048xf32> loc(#loc5262)
    %8111 = stablehlo.broadcast_in_dim %8081, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8112 = stablehlo.broadcast_in_dim %8111, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8113 = stablehlo.multiply %8112, %arg245 : tensor<1x2048xf32> loc(#loc5262)
    %8114 = stablehlo.add %8110, %8113 : tensor<1x2048xf32> loc(#loc5263)
    %8115 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8116 = stablehlo.broadcast_in_dim %8115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8117 = stablehlo.subtract %8116, %8104 : tensor<1xf32> loc(#loc5260)
    %8118 = stablehlo.multiply %4460, %4460 : tensor<1x2048xf32> loc(#loc5264)
    %8119 = stablehlo.broadcast_in_dim %8117, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8120 = stablehlo.broadcast_in_dim %8119, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8121 = stablehlo.multiply %8120, %8118 : tensor<1x2048xf32> loc(#loc5262)
    %8122 = stablehlo.broadcast_in_dim %8104, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8123 = stablehlo.broadcast_in_dim %8122, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8124 = stablehlo.multiply %8123, %arg421 : tensor<1x2048xf32> loc(#loc5262)
    %8125 = stablehlo.add %8121, %8124 : tensor<1x2048xf32> loc(#loc5263)
    %8126 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8127 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8128 = stablehlo.broadcast_in_dim %8127, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8129 = stablehlo.add %8126, %8128 : tensor<1xf32> loc(#loc5259)
    %8130 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8131 = stablehlo.broadcast_in_dim %8130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8132 = stablehlo.subtract %8129, %8131 : tensor<1xf32> loc(#loc5260)
    %8133 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8134 = stablehlo.broadcast_in_dim %8133, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8135 = stablehlo.power %8134, %8132 : tensor<1xf32> loc(#loc5261)
    %8136 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8137 = stablehlo.broadcast_in_dim %8136, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8138 = stablehlo.subtract %8137, %8135 : tensor<1xf32> loc(#loc5260)
    %8139 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8140 = stablehlo.broadcast_in_dim %8139, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8141 = stablehlo.multiply %8140, %8138 : tensor<1xf32> loc(#loc5236)
    %8142 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8143 = stablehlo.broadcast_in_dim %8142, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8144 = stablehlo.power %8143, %8129 : tensor<1xf32> loc(#loc5261)
    %8145 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8146 = stablehlo.broadcast_in_dim %8145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8147 = stablehlo.subtract %8146, %8144 : tensor<1xf32> loc(#loc5260)
    %8148 = stablehlo.divide %8141, %8147 : tensor<1xf32> loc(#loc5253)
    %8149 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8150 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8151 = stablehlo.broadcast_in_dim %8150, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8152 = stablehlo.add %8149, %8151 : tensor<1xf32> loc(#loc5259)
    %8153 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8154 = stablehlo.broadcast_in_dim %8153, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8155 = stablehlo.subtract %8152, %8154 : tensor<1xf32> loc(#loc5260)
    %8156 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8157 = stablehlo.broadcast_in_dim %8156, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8158 = stablehlo.power %8157, %8155 : tensor<1xf32> loc(#loc5261)
    %8159 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8160 = stablehlo.broadcast_in_dim %8159, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8161 = stablehlo.subtract %8160, %8158 : tensor<1xf32> loc(#loc5260)
    %8162 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8163 = stablehlo.broadcast_in_dim %8162, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8164 = stablehlo.multiply %8163, %8161 : tensor<1xf32> loc(#loc5236)
    %8165 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8166 = stablehlo.broadcast_in_dim %8165, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8167 = stablehlo.power %8166, %8152 : tensor<1xf32> loc(#loc5261)
    %8168 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8169 = stablehlo.broadcast_in_dim %8168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8170 = stablehlo.subtract %8169, %8167 : tensor<1xf32> loc(#loc5260)
    %8171 = stablehlo.divide %8164, %8170 : tensor<1xf32> loc(#loc5253)
    %8172 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8173 = stablehlo.broadcast_in_dim %8172, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8174 = stablehlo.subtract %8173, %8148 : tensor<1xf32> loc(#loc5260)
    %8175 = stablehlo.broadcast_in_dim %8174, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8176 = stablehlo.broadcast_in_dim %8175, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8177 = stablehlo.multiply %8176, %4464 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8178 = stablehlo.broadcast_in_dim %8148, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8179 = stablehlo.broadcast_in_dim %8178, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8180 = stablehlo.multiply %8179, %arg246 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8181 = stablehlo.add %8177, %8180 : tensor<1x8192x2048xf32> loc(#loc5263)
    %8182 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8183 = stablehlo.broadcast_in_dim %8182, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8184 = stablehlo.subtract %8183, %8171 : tensor<1xf32> loc(#loc5260)
    %8185 = stablehlo.multiply %4464, %4464 : tensor<1x8192x2048xf32> loc(#loc5264)
    %8186 = stablehlo.broadcast_in_dim %8184, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8187 = stablehlo.broadcast_in_dim %8186, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8188 = stablehlo.multiply %8187, %8185 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8189 = stablehlo.broadcast_in_dim %8171, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8190 = stablehlo.broadcast_in_dim %8189, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8191 = stablehlo.multiply %8190, %arg422 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8192 = stablehlo.add %8188, %8191 : tensor<1x8192x2048xf32> loc(#loc5263)
    %8193 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8194 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8195 = stablehlo.broadcast_in_dim %8194, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8196 = stablehlo.add %8193, %8195 : tensor<1xf32> loc(#loc5259)
    %8197 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8198 = stablehlo.broadcast_in_dim %8197, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8199 = stablehlo.subtract %8196, %8198 : tensor<1xf32> loc(#loc5260)
    %8200 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8201 = stablehlo.broadcast_in_dim %8200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8202 = stablehlo.power %8201, %8199 : tensor<1xf32> loc(#loc5261)
    %8203 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8204 = stablehlo.broadcast_in_dim %8203, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8205 = stablehlo.subtract %8204, %8202 : tensor<1xf32> loc(#loc5260)
    %8206 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8207 = stablehlo.broadcast_in_dim %8206, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8208 = stablehlo.multiply %8207, %8205 : tensor<1xf32> loc(#loc5236)
    %8209 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8210 = stablehlo.broadcast_in_dim %8209, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8211 = stablehlo.power %8210, %8196 : tensor<1xf32> loc(#loc5261)
    %8212 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8213 = stablehlo.broadcast_in_dim %8212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8214 = stablehlo.subtract %8213, %8211 : tensor<1xf32> loc(#loc5260)
    %8215 = stablehlo.divide %8208, %8214 : tensor<1xf32> loc(#loc5253)
    %8216 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8217 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8218 = stablehlo.broadcast_in_dim %8217, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8219 = stablehlo.add %8216, %8218 : tensor<1xf32> loc(#loc5259)
    %8220 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8221 = stablehlo.broadcast_in_dim %8220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8222 = stablehlo.subtract %8219, %8221 : tensor<1xf32> loc(#loc5260)
    %8223 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8224 = stablehlo.broadcast_in_dim %8223, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8225 = stablehlo.power %8224, %8222 : tensor<1xf32> loc(#loc5261)
    %8226 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8227 = stablehlo.broadcast_in_dim %8226, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8228 = stablehlo.subtract %8227, %8225 : tensor<1xf32> loc(#loc5260)
    %8229 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8230 = stablehlo.broadcast_in_dim %8229, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8231 = stablehlo.multiply %8230, %8228 : tensor<1xf32> loc(#loc5236)
    %8232 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8233 = stablehlo.broadcast_in_dim %8232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8234 = stablehlo.power %8233, %8219 : tensor<1xf32> loc(#loc5261)
    %8235 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8236 = stablehlo.broadcast_in_dim %8235, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8237 = stablehlo.subtract %8236, %8234 : tensor<1xf32> loc(#loc5260)
    %8238 = stablehlo.divide %8231, %8237 : tensor<1xf32> loc(#loc5253)
    %8239 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8240 = stablehlo.broadcast_in_dim %8239, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8241 = stablehlo.subtract %8240, %8215 : tensor<1xf32> loc(#loc5260)
    %8242 = stablehlo.broadcast_in_dim %8241, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8243 = stablehlo.broadcast_in_dim %8242, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8244 = stablehlo.multiply %8243, %4468 : tensor<1x2048xf32> loc(#loc5262)
    %8245 = stablehlo.broadcast_in_dim %8215, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8246 = stablehlo.broadcast_in_dim %8245, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8247 = stablehlo.multiply %8246, %arg247 : tensor<1x2048xf32> loc(#loc5262)
    %8248 = stablehlo.add %8244, %8247 : tensor<1x2048xf32> loc(#loc5263)
    %8249 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8250 = stablehlo.broadcast_in_dim %8249, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8251 = stablehlo.subtract %8250, %8238 : tensor<1xf32> loc(#loc5260)
    %8252 = stablehlo.multiply %4468, %4468 : tensor<1x2048xf32> loc(#loc5264)
    %8253 = stablehlo.broadcast_in_dim %8251, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8254 = stablehlo.broadcast_in_dim %8253, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8255 = stablehlo.multiply %8254, %8252 : tensor<1x2048xf32> loc(#loc5262)
    %8256 = stablehlo.broadcast_in_dim %8238, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8257 = stablehlo.broadcast_in_dim %8256, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8258 = stablehlo.multiply %8257, %arg423 : tensor<1x2048xf32> loc(#loc5262)
    %8259 = stablehlo.add %8255, %8258 : tensor<1x2048xf32> loc(#loc5263)
    %8260 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8261 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8262 = stablehlo.broadcast_in_dim %8261, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8263 = stablehlo.add %8260, %8262 : tensor<1xf32> loc(#loc5259)
    %8264 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8265 = stablehlo.broadcast_in_dim %8264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8266 = stablehlo.subtract %8263, %8265 : tensor<1xf32> loc(#loc5260)
    %8267 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8268 = stablehlo.broadcast_in_dim %8267, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8269 = stablehlo.power %8268, %8266 : tensor<1xf32> loc(#loc5261)
    %8270 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8271 = stablehlo.broadcast_in_dim %8270, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8272 = stablehlo.subtract %8271, %8269 : tensor<1xf32> loc(#loc5260)
    %8273 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8274 = stablehlo.broadcast_in_dim %8273, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8275 = stablehlo.multiply %8274, %8272 : tensor<1xf32> loc(#loc5236)
    %8276 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8277 = stablehlo.broadcast_in_dim %8276, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8278 = stablehlo.power %8277, %8263 : tensor<1xf32> loc(#loc5261)
    %8279 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8280 = stablehlo.broadcast_in_dim %8279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8281 = stablehlo.subtract %8280, %8278 : tensor<1xf32> loc(#loc5260)
    %8282 = stablehlo.divide %8275, %8281 : tensor<1xf32> loc(#loc5253)
    %8283 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8284 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8285 = stablehlo.broadcast_in_dim %8284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8286 = stablehlo.add %8283, %8285 : tensor<1xf32> loc(#loc5259)
    %8287 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8288 = stablehlo.broadcast_in_dim %8287, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8289 = stablehlo.subtract %8286, %8288 : tensor<1xf32> loc(#loc5260)
    %8290 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8291 = stablehlo.broadcast_in_dim %8290, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8292 = stablehlo.power %8291, %8289 : tensor<1xf32> loc(#loc5261)
    %8293 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8294 = stablehlo.broadcast_in_dim %8293, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8295 = stablehlo.subtract %8294, %8292 : tensor<1xf32> loc(#loc5260)
    %8296 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8297 = stablehlo.broadcast_in_dim %8296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8298 = stablehlo.multiply %8297, %8295 : tensor<1xf32> loc(#loc5236)
    %8299 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8300 = stablehlo.broadcast_in_dim %8299, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8301 = stablehlo.power %8300, %8286 : tensor<1xf32> loc(#loc5261)
    %8302 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8303 = stablehlo.broadcast_in_dim %8302, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8304 = stablehlo.subtract %8303, %8301 : tensor<1xf32> loc(#loc5260)
    %8305 = stablehlo.divide %8298, %8304 : tensor<1xf32> loc(#loc5253)
    %8306 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8307 = stablehlo.broadcast_in_dim %8306, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8308 = stablehlo.subtract %8307, %8282 : tensor<1xf32> loc(#loc5260)
    %8309 = stablehlo.broadcast_in_dim %8308, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8310 = stablehlo.broadcast_in_dim %8309, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8311 = stablehlo.multiply %8310, %4472 : tensor<1x2048xf32> loc(#loc5262)
    %8312 = stablehlo.broadcast_in_dim %8282, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8313 = stablehlo.broadcast_in_dim %8312, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8314 = stablehlo.multiply %8313, %arg248 : tensor<1x2048xf32> loc(#loc5262)
    %8315 = stablehlo.add %8311, %8314 : tensor<1x2048xf32> loc(#loc5263)
    %8316 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8317 = stablehlo.broadcast_in_dim %8316, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8318 = stablehlo.subtract %8317, %8305 : tensor<1xf32> loc(#loc5260)
    %8319 = stablehlo.multiply %4472, %4472 : tensor<1x2048xf32> loc(#loc5264)
    %8320 = stablehlo.broadcast_in_dim %8318, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8321 = stablehlo.broadcast_in_dim %8320, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8322 = stablehlo.multiply %8321, %8319 : tensor<1x2048xf32> loc(#loc5262)
    %8323 = stablehlo.broadcast_in_dim %8305, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8324 = stablehlo.broadcast_in_dim %8323, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8325 = stablehlo.multiply %8324, %arg424 : tensor<1x2048xf32> loc(#loc5262)
    %8326 = stablehlo.add %8322, %8325 : tensor<1x2048xf32> loc(#loc5263)
    %8327 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8328 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8329 = stablehlo.broadcast_in_dim %8328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8330 = stablehlo.add %8327, %8329 : tensor<1xf32> loc(#loc5259)
    %8331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8332 = stablehlo.broadcast_in_dim %8331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8333 = stablehlo.subtract %8330, %8332 : tensor<1xf32> loc(#loc5260)
    %8334 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8335 = stablehlo.broadcast_in_dim %8334, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8336 = stablehlo.power %8335, %8333 : tensor<1xf32> loc(#loc5261)
    %8337 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8338 = stablehlo.broadcast_in_dim %8337, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8339 = stablehlo.subtract %8338, %8336 : tensor<1xf32> loc(#loc5260)
    %8340 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8341 = stablehlo.broadcast_in_dim %8340, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8342 = stablehlo.multiply %8341, %8339 : tensor<1xf32> loc(#loc5236)
    %8343 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8344 = stablehlo.broadcast_in_dim %8343, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8345 = stablehlo.power %8344, %8330 : tensor<1xf32> loc(#loc5261)
    %8346 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8347 = stablehlo.broadcast_in_dim %8346, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8348 = stablehlo.subtract %8347, %8345 : tensor<1xf32> loc(#loc5260)
    %8349 = stablehlo.divide %8342, %8348 : tensor<1xf32> loc(#loc5253)
    %8350 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8352 = stablehlo.broadcast_in_dim %8351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8353 = stablehlo.add %8350, %8352 : tensor<1xf32> loc(#loc5259)
    %8354 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8355 = stablehlo.broadcast_in_dim %8354, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8356 = stablehlo.subtract %8353, %8355 : tensor<1xf32> loc(#loc5260)
    %8357 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8358 = stablehlo.broadcast_in_dim %8357, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8359 = stablehlo.power %8358, %8356 : tensor<1xf32> loc(#loc5261)
    %8360 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8361 = stablehlo.broadcast_in_dim %8360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8362 = stablehlo.subtract %8361, %8359 : tensor<1xf32> loc(#loc5260)
    %8363 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8364 = stablehlo.broadcast_in_dim %8363, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8365 = stablehlo.multiply %8364, %8362 : tensor<1xf32> loc(#loc5236)
    %8366 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8367 = stablehlo.broadcast_in_dim %8366, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8368 = stablehlo.power %8367, %8353 : tensor<1xf32> loc(#loc5261)
    %8369 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8370 = stablehlo.broadcast_in_dim %8369, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8371 = stablehlo.subtract %8370, %8368 : tensor<1xf32> loc(#loc5260)
    %8372 = stablehlo.divide %8365, %8371 : tensor<1xf32> loc(#loc5253)
    %8373 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8374 = stablehlo.broadcast_in_dim %8373, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8375 = stablehlo.subtract %8374, %8349 : tensor<1xf32> loc(#loc5260)
    %8376 = stablehlo.broadcast_in_dim %8375, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8377 = stablehlo.broadcast_in_dim %8376, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8378 = stablehlo.multiply %8377, %4476 : tensor<1x2048xf32> loc(#loc5262)
    %8379 = stablehlo.broadcast_in_dim %8349, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8380 = stablehlo.broadcast_in_dim %8379, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8381 = stablehlo.multiply %8380, %arg249 : tensor<1x2048xf32> loc(#loc5262)
    %8382 = stablehlo.add %8378, %8381 : tensor<1x2048xf32> loc(#loc5263)
    %8383 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8384 = stablehlo.broadcast_in_dim %8383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8385 = stablehlo.subtract %8384, %8372 : tensor<1xf32> loc(#loc5260)
    %8386 = stablehlo.multiply %4476, %4476 : tensor<1x2048xf32> loc(#loc5264)
    %8387 = stablehlo.broadcast_in_dim %8385, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8388 = stablehlo.broadcast_in_dim %8387, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8389 = stablehlo.multiply %8388, %8386 : tensor<1x2048xf32> loc(#loc5262)
    %8390 = stablehlo.broadcast_in_dim %8372, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8391 = stablehlo.broadcast_in_dim %8390, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8392 = stablehlo.multiply %8391, %arg425 : tensor<1x2048xf32> loc(#loc5262)
    %8393 = stablehlo.add %8389, %8392 : tensor<1x2048xf32> loc(#loc5263)
    %8394 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8395 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8396 = stablehlo.broadcast_in_dim %8395, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8397 = stablehlo.add %8394, %8396 : tensor<1xf32> loc(#loc5259)
    %8398 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8399 = stablehlo.broadcast_in_dim %8398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8400 = stablehlo.subtract %8397, %8399 : tensor<1xf32> loc(#loc5260)
    %8401 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8402 = stablehlo.broadcast_in_dim %8401, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8403 = stablehlo.power %8402, %8400 : tensor<1xf32> loc(#loc5261)
    %8404 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8405 = stablehlo.broadcast_in_dim %8404, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8406 = stablehlo.subtract %8405, %8403 : tensor<1xf32> loc(#loc5260)
    %8407 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8408 = stablehlo.broadcast_in_dim %8407, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8409 = stablehlo.multiply %8408, %8406 : tensor<1xf32> loc(#loc5236)
    %8410 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8411 = stablehlo.broadcast_in_dim %8410, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8412 = stablehlo.power %8411, %8397 : tensor<1xf32> loc(#loc5261)
    %8413 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8414 = stablehlo.broadcast_in_dim %8413, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8415 = stablehlo.subtract %8414, %8412 : tensor<1xf32> loc(#loc5260)
    %8416 = stablehlo.divide %8409, %8415 : tensor<1xf32> loc(#loc5253)
    %8417 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8418 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8419 = stablehlo.broadcast_in_dim %8418, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8420 = stablehlo.add %8417, %8419 : tensor<1xf32> loc(#loc5259)
    %8421 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8422 = stablehlo.broadcast_in_dim %8421, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8423 = stablehlo.subtract %8420, %8422 : tensor<1xf32> loc(#loc5260)
    %8424 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8425 = stablehlo.broadcast_in_dim %8424, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8426 = stablehlo.power %8425, %8423 : tensor<1xf32> loc(#loc5261)
    %8427 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8428 = stablehlo.broadcast_in_dim %8427, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8429 = stablehlo.subtract %8428, %8426 : tensor<1xf32> loc(#loc5260)
    %8430 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8431 = stablehlo.broadcast_in_dim %8430, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8432 = stablehlo.multiply %8431, %8429 : tensor<1xf32> loc(#loc5236)
    %8433 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8434 = stablehlo.broadcast_in_dim %8433, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8435 = stablehlo.power %8434, %8420 : tensor<1xf32> loc(#loc5261)
    %8436 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8437 = stablehlo.broadcast_in_dim %8436, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8438 = stablehlo.subtract %8437, %8435 : tensor<1xf32> loc(#loc5260)
    %8439 = stablehlo.divide %8432, %8438 : tensor<1xf32> loc(#loc5253)
    %8440 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8441 = stablehlo.broadcast_in_dim %8440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8442 = stablehlo.subtract %8441, %8416 : tensor<1xf32> loc(#loc5260)
    %8443 = stablehlo.broadcast_in_dim %8442, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8444 = stablehlo.broadcast_in_dim %8443, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8445 = stablehlo.multiply %8444, %4480 : tensor<1x2048xf32> loc(#loc5262)
    %8446 = stablehlo.broadcast_in_dim %8416, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8447 = stablehlo.broadcast_in_dim %8446, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8448 = stablehlo.multiply %8447, %arg250 : tensor<1x2048xf32> loc(#loc5262)
    %8449 = stablehlo.add %8445, %8448 : tensor<1x2048xf32> loc(#loc5263)
    %8450 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8451 = stablehlo.broadcast_in_dim %8450, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8452 = stablehlo.subtract %8451, %8439 : tensor<1xf32> loc(#loc5260)
    %8453 = stablehlo.multiply %4480, %4480 : tensor<1x2048xf32> loc(#loc5264)
    %8454 = stablehlo.broadcast_in_dim %8452, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8455 = stablehlo.broadcast_in_dim %8454, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8456 = stablehlo.multiply %8455, %8453 : tensor<1x2048xf32> loc(#loc5262)
    %8457 = stablehlo.broadcast_in_dim %8439, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8458 = stablehlo.broadcast_in_dim %8457, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8459 = stablehlo.multiply %8458, %arg426 : tensor<1x2048xf32> loc(#loc5262)
    %8460 = stablehlo.add %8456, %8459 : tensor<1x2048xf32> loc(#loc5263)
    %8461 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8462 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8463 = stablehlo.broadcast_in_dim %8462, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8464 = stablehlo.add %8461, %8463 : tensor<1xf32> loc(#loc5259)
    %8465 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8466 = stablehlo.broadcast_in_dim %8465, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8467 = stablehlo.subtract %8464, %8466 : tensor<1xf32> loc(#loc5260)
    %8468 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8469 = stablehlo.broadcast_in_dim %8468, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8470 = stablehlo.power %8469, %8467 : tensor<1xf32> loc(#loc5261)
    %8471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8472 = stablehlo.broadcast_in_dim %8471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8473 = stablehlo.subtract %8472, %8470 : tensor<1xf32> loc(#loc5260)
    %8474 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8475 = stablehlo.broadcast_in_dim %8474, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8476 = stablehlo.multiply %8475, %8473 : tensor<1xf32> loc(#loc5236)
    %8477 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8478 = stablehlo.broadcast_in_dim %8477, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8479 = stablehlo.power %8478, %8464 : tensor<1xf32> loc(#loc5261)
    %8480 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8481 = stablehlo.broadcast_in_dim %8480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8482 = stablehlo.subtract %8481, %8479 : tensor<1xf32> loc(#loc5260)
    %8483 = stablehlo.divide %8476, %8482 : tensor<1xf32> loc(#loc5253)
    %8484 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8485 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8486 = stablehlo.broadcast_in_dim %8485, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8487 = stablehlo.add %8484, %8486 : tensor<1xf32> loc(#loc5259)
    %8488 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8489 = stablehlo.broadcast_in_dim %8488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8490 = stablehlo.subtract %8487, %8489 : tensor<1xf32> loc(#loc5260)
    %8491 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8492 = stablehlo.broadcast_in_dim %8491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8493 = stablehlo.power %8492, %8490 : tensor<1xf32> loc(#loc5261)
    %8494 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8495 = stablehlo.broadcast_in_dim %8494, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8496 = stablehlo.subtract %8495, %8493 : tensor<1xf32> loc(#loc5260)
    %8497 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8498 = stablehlo.broadcast_in_dim %8497, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8499 = stablehlo.multiply %8498, %8496 : tensor<1xf32> loc(#loc5236)
    %8500 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8501 = stablehlo.broadcast_in_dim %8500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8502 = stablehlo.power %8501, %8487 : tensor<1xf32> loc(#loc5261)
    %8503 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8504 = stablehlo.broadcast_in_dim %8503, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8505 = stablehlo.subtract %8504, %8502 : tensor<1xf32> loc(#loc5260)
    %8506 = stablehlo.divide %8499, %8505 : tensor<1xf32> loc(#loc5253)
    %8507 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8508 = stablehlo.broadcast_in_dim %8507, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8509 = stablehlo.subtract %8508, %8483 : tensor<1xf32> loc(#loc5260)
    %8510 = stablehlo.broadcast_in_dim %8509, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %8511 = stablehlo.broadcast_in_dim %8510, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8512 = stablehlo.multiply %8511, %4484 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8513 = stablehlo.broadcast_in_dim %8483, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %8514 = stablehlo.broadcast_in_dim %8513, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8515 = stablehlo.multiply %8514, %arg251 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8516 = stablehlo.add %8512, %8515 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %8517 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8518 = stablehlo.broadcast_in_dim %8517, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8519 = stablehlo.subtract %8518, %8506 : tensor<1xf32> loc(#loc5260)
    %8520 = stablehlo.multiply %4484, %4484 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %8521 = stablehlo.broadcast_in_dim %8519, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %8522 = stablehlo.broadcast_in_dim %8521, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8523 = stablehlo.multiply %8522, %8520 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8524 = stablehlo.broadcast_in_dim %8506, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %8525 = stablehlo.broadcast_in_dim %8524, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8526 = stablehlo.multiply %8525, %arg427 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %8527 = stablehlo.add %8523, %8526 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %8528 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8529 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8530 = stablehlo.broadcast_in_dim %8529, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8531 = stablehlo.add %8528, %8530 : tensor<1xf32> loc(#loc5259)
    %8532 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8533 = stablehlo.broadcast_in_dim %8532, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8534 = stablehlo.subtract %8531, %8533 : tensor<1xf32> loc(#loc5260)
    %8535 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8536 = stablehlo.broadcast_in_dim %8535, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8537 = stablehlo.power %8536, %8534 : tensor<1xf32> loc(#loc5261)
    %8538 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8539 = stablehlo.broadcast_in_dim %8538, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8540 = stablehlo.subtract %8539, %8537 : tensor<1xf32> loc(#loc5260)
    %8541 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8542 = stablehlo.broadcast_in_dim %8541, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8543 = stablehlo.multiply %8542, %8540 : tensor<1xf32> loc(#loc5236)
    %8544 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8545 = stablehlo.broadcast_in_dim %8544, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8546 = stablehlo.power %8545, %8531 : tensor<1xf32> loc(#loc5261)
    %8547 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8548 = stablehlo.broadcast_in_dim %8547, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8549 = stablehlo.subtract %8548, %8546 : tensor<1xf32> loc(#loc5260)
    %8550 = stablehlo.divide %8543, %8549 : tensor<1xf32> loc(#loc5253)
    %8551 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8552 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8553 = stablehlo.broadcast_in_dim %8552, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8554 = stablehlo.add %8551, %8553 : tensor<1xf32> loc(#loc5259)
    %8555 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8556 = stablehlo.broadcast_in_dim %8555, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8557 = stablehlo.subtract %8554, %8556 : tensor<1xf32> loc(#loc5260)
    %8558 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8559 = stablehlo.broadcast_in_dim %8558, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8560 = stablehlo.power %8559, %8557 : tensor<1xf32> loc(#loc5261)
    %8561 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8562 = stablehlo.broadcast_in_dim %8561, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8563 = stablehlo.subtract %8562, %8560 : tensor<1xf32> loc(#loc5260)
    %8564 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8565 = stablehlo.broadcast_in_dim %8564, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8566 = stablehlo.multiply %8565, %8563 : tensor<1xf32> loc(#loc5236)
    %8567 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8568 = stablehlo.broadcast_in_dim %8567, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8569 = stablehlo.power %8568, %8554 : tensor<1xf32> loc(#loc5261)
    %8570 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8571 = stablehlo.broadcast_in_dim %8570, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8572 = stablehlo.subtract %8571, %8569 : tensor<1xf32> loc(#loc5260)
    %8573 = stablehlo.divide %8566, %8572 : tensor<1xf32> loc(#loc5253)
    %8574 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8575 = stablehlo.broadcast_in_dim %8574, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8576 = stablehlo.subtract %8575, %8550 : tensor<1xf32> loc(#loc5260)
    %8577 = stablehlo.broadcast_in_dim %8576, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8578 = stablehlo.broadcast_in_dim %8577, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %8579 = stablehlo.multiply %8578, %4488 : tensor<1x64xf32> loc(#loc5262)
    %8580 = stablehlo.broadcast_in_dim %8550, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8581 = stablehlo.broadcast_in_dim %8580, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %8582 = stablehlo.multiply %8581, %arg252 : tensor<1x64xf32> loc(#loc5262)
    %8583 = stablehlo.add %8579, %8582 : tensor<1x64xf32> loc(#loc5263)
    %8584 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8585 = stablehlo.broadcast_in_dim %8584, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8586 = stablehlo.subtract %8585, %8573 : tensor<1xf32> loc(#loc5260)
    %8587 = stablehlo.multiply %4488, %4488 : tensor<1x64xf32> loc(#loc5264)
    %8588 = stablehlo.broadcast_in_dim %8586, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8589 = stablehlo.broadcast_in_dim %8588, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %8590 = stablehlo.multiply %8589, %8587 : tensor<1x64xf32> loc(#loc5262)
    %8591 = stablehlo.broadcast_in_dim %8573, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8592 = stablehlo.broadcast_in_dim %8591, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %8593 = stablehlo.multiply %8592, %arg428 : tensor<1x64xf32> loc(#loc5262)
    %8594 = stablehlo.add %8590, %8593 : tensor<1x64xf32> loc(#loc5263)
    %8595 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8596 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8597 = stablehlo.broadcast_in_dim %8596, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8598 = stablehlo.add %8595, %8597 : tensor<1xf32> loc(#loc5259)
    %8599 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8600 = stablehlo.broadcast_in_dim %8599, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8601 = stablehlo.subtract %8598, %8600 : tensor<1xf32> loc(#loc5260)
    %8602 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8603 = stablehlo.broadcast_in_dim %8602, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8604 = stablehlo.power %8603, %8601 : tensor<1xf32> loc(#loc5261)
    %8605 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8606 = stablehlo.broadcast_in_dim %8605, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8607 = stablehlo.subtract %8606, %8604 : tensor<1xf32> loc(#loc5260)
    %8608 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8609 = stablehlo.broadcast_in_dim %8608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8610 = stablehlo.multiply %8609, %8607 : tensor<1xf32> loc(#loc5236)
    %8611 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8612 = stablehlo.broadcast_in_dim %8611, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8613 = stablehlo.power %8612, %8598 : tensor<1xf32> loc(#loc5261)
    %8614 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8615 = stablehlo.broadcast_in_dim %8614, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8616 = stablehlo.subtract %8615, %8613 : tensor<1xf32> loc(#loc5260)
    %8617 = stablehlo.divide %8610, %8616 : tensor<1xf32> loc(#loc5253)
    %8618 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8619 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8620 = stablehlo.broadcast_in_dim %8619, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8621 = stablehlo.add %8618, %8620 : tensor<1xf32> loc(#loc5259)
    %8622 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8623 = stablehlo.broadcast_in_dim %8622, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8624 = stablehlo.subtract %8621, %8623 : tensor<1xf32> loc(#loc5260)
    %8625 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8626 = stablehlo.broadcast_in_dim %8625, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8627 = stablehlo.power %8626, %8624 : tensor<1xf32> loc(#loc5261)
    %8628 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8629 = stablehlo.broadcast_in_dim %8628, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8630 = stablehlo.subtract %8629, %8627 : tensor<1xf32> loc(#loc5260)
    %8631 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8632 = stablehlo.broadcast_in_dim %8631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8633 = stablehlo.multiply %8632, %8630 : tensor<1xf32> loc(#loc5236)
    %8634 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8635 = stablehlo.broadcast_in_dim %8634, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8636 = stablehlo.power %8635, %8621 : tensor<1xf32> loc(#loc5261)
    %8637 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8638 = stablehlo.broadcast_in_dim %8637, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8639 = stablehlo.subtract %8638, %8636 : tensor<1xf32> loc(#loc5260)
    %8640 = stablehlo.divide %8633, %8639 : tensor<1xf32> loc(#loc5253)
    %8641 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8642 = stablehlo.broadcast_in_dim %8641, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8643 = stablehlo.subtract %8642, %8617 : tensor<1xf32> loc(#loc5260)
    %8644 = stablehlo.broadcast_in_dim %8643, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %8645 = stablehlo.broadcast_in_dim %8644, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %8646 = stablehlo.multiply %8645, %4492 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %8647 = stablehlo.broadcast_in_dim %8617, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %8648 = stablehlo.broadcast_in_dim %8647, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %8649 = stablehlo.multiply %8648, %arg253 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %8650 = stablehlo.add %8646, %8649 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %8651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8652 = stablehlo.broadcast_in_dim %8651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8653 = stablehlo.subtract %8652, %8640 : tensor<1xf32> loc(#loc5260)
    %8654 = stablehlo.multiply %4492, %4492 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %8655 = stablehlo.broadcast_in_dim %8653, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %8656 = stablehlo.broadcast_in_dim %8655, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %8657 = stablehlo.multiply %8656, %8654 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %8658 = stablehlo.broadcast_in_dim %8640, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %8659 = stablehlo.broadcast_in_dim %8658, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %8660 = stablehlo.multiply %8659, %arg429 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %8661 = stablehlo.add %8657, %8660 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %8662 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8663 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8664 = stablehlo.broadcast_in_dim %8663, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8665 = stablehlo.add %8662, %8664 : tensor<1xf32> loc(#loc5259)
    %8666 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8667 = stablehlo.broadcast_in_dim %8666, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8668 = stablehlo.subtract %8665, %8667 : tensor<1xf32> loc(#loc5260)
    %8669 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8670 = stablehlo.broadcast_in_dim %8669, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8671 = stablehlo.power %8670, %8668 : tensor<1xf32> loc(#loc5261)
    %8672 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8673 = stablehlo.broadcast_in_dim %8672, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8674 = stablehlo.subtract %8673, %8671 : tensor<1xf32> loc(#loc5260)
    %8675 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8676 = stablehlo.broadcast_in_dim %8675, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8677 = stablehlo.multiply %8676, %8674 : tensor<1xf32> loc(#loc5236)
    %8678 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8679 = stablehlo.broadcast_in_dim %8678, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8680 = stablehlo.power %8679, %8665 : tensor<1xf32> loc(#loc5261)
    %8681 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8682 = stablehlo.broadcast_in_dim %8681, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8683 = stablehlo.subtract %8682, %8680 : tensor<1xf32> loc(#loc5260)
    %8684 = stablehlo.divide %8677, %8683 : tensor<1xf32> loc(#loc5253)
    %8685 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8686 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8687 = stablehlo.broadcast_in_dim %8686, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8688 = stablehlo.add %8685, %8687 : tensor<1xf32> loc(#loc5259)
    %8689 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8690 = stablehlo.broadcast_in_dim %8689, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8691 = stablehlo.subtract %8688, %8690 : tensor<1xf32> loc(#loc5260)
    %8692 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8693 = stablehlo.broadcast_in_dim %8692, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8694 = stablehlo.power %8693, %8691 : tensor<1xf32> loc(#loc5261)
    %8695 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8696 = stablehlo.broadcast_in_dim %8695, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8697 = stablehlo.subtract %8696, %8694 : tensor<1xf32> loc(#loc5260)
    %8698 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8699 = stablehlo.broadcast_in_dim %8698, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8700 = stablehlo.multiply %8699, %8697 : tensor<1xf32> loc(#loc5236)
    %8701 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8702 = stablehlo.broadcast_in_dim %8701, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8703 = stablehlo.power %8702, %8688 : tensor<1xf32> loc(#loc5261)
    %8704 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8705 = stablehlo.broadcast_in_dim %8704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8706 = stablehlo.subtract %8705, %8703 : tensor<1xf32> loc(#loc5260)
    %8707 = stablehlo.divide %8700, %8706 : tensor<1xf32> loc(#loc5253)
    %8708 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8709 = stablehlo.broadcast_in_dim %8708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8710 = stablehlo.subtract %8709, %8684 : tensor<1xf32> loc(#loc5260)
    %8711 = stablehlo.broadcast_in_dim %8710, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8712 = stablehlo.broadcast_in_dim %8711, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %8713 = stablehlo.multiply %8712, %4496 : tensor<1x8192xf32> loc(#loc5262)
    %8714 = stablehlo.broadcast_in_dim %8684, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8715 = stablehlo.broadcast_in_dim %8714, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %8716 = stablehlo.multiply %8715, %arg254 : tensor<1x8192xf32> loc(#loc5262)
    %8717 = stablehlo.add %8713, %8716 : tensor<1x8192xf32> loc(#loc5263)
    %8718 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8719 = stablehlo.broadcast_in_dim %8718, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8720 = stablehlo.subtract %8719, %8707 : tensor<1xf32> loc(#loc5260)
    %8721 = stablehlo.multiply %4496, %4496 : tensor<1x8192xf32> loc(#loc5264)
    %8722 = stablehlo.broadcast_in_dim %8720, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8723 = stablehlo.broadcast_in_dim %8722, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %8724 = stablehlo.multiply %8723, %8721 : tensor<1x8192xf32> loc(#loc5262)
    %8725 = stablehlo.broadcast_in_dim %8707, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8726 = stablehlo.broadcast_in_dim %8725, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %8727 = stablehlo.multiply %8726, %arg430 : tensor<1x8192xf32> loc(#loc5262)
    %8728 = stablehlo.add %8724, %8727 : tensor<1x8192xf32> loc(#loc5263)
    %8729 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8730 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8731 = stablehlo.broadcast_in_dim %8730, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8732 = stablehlo.add %8729, %8731 : tensor<1xf32> loc(#loc5259)
    %8733 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8734 = stablehlo.broadcast_in_dim %8733, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8735 = stablehlo.subtract %8732, %8734 : tensor<1xf32> loc(#loc5260)
    %8736 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8737 = stablehlo.broadcast_in_dim %8736, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8738 = stablehlo.power %8737, %8735 : tensor<1xf32> loc(#loc5261)
    %8739 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8740 = stablehlo.broadcast_in_dim %8739, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8741 = stablehlo.subtract %8740, %8738 : tensor<1xf32> loc(#loc5260)
    %8742 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8743 = stablehlo.broadcast_in_dim %8742, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8744 = stablehlo.multiply %8743, %8741 : tensor<1xf32> loc(#loc5236)
    %8745 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8746 = stablehlo.broadcast_in_dim %8745, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8747 = stablehlo.power %8746, %8732 : tensor<1xf32> loc(#loc5261)
    %8748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8749 = stablehlo.broadcast_in_dim %8748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8750 = stablehlo.subtract %8749, %8747 : tensor<1xf32> loc(#loc5260)
    %8751 = stablehlo.divide %8744, %8750 : tensor<1xf32> loc(#loc5253)
    %8752 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8753 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8754 = stablehlo.broadcast_in_dim %8753, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8755 = stablehlo.add %8752, %8754 : tensor<1xf32> loc(#loc5259)
    %8756 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8757 = stablehlo.broadcast_in_dim %8756, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8758 = stablehlo.subtract %8755, %8757 : tensor<1xf32> loc(#loc5260)
    %8759 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8760 = stablehlo.broadcast_in_dim %8759, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8761 = stablehlo.power %8760, %8758 : tensor<1xf32> loc(#loc5261)
    %8762 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8763 = stablehlo.broadcast_in_dim %8762, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8764 = stablehlo.subtract %8763, %8761 : tensor<1xf32> loc(#loc5260)
    %8765 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8766 = stablehlo.broadcast_in_dim %8765, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8767 = stablehlo.multiply %8766, %8764 : tensor<1xf32> loc(#loc5236)
    %8768 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8769 = stablehlo.broadcast_in_dim %8768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8770 = stablehlo.power %8769, %8755 : tensor<1xf32> loc(#loc5261)
    %8771 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8772 = stablehlo.broadcast_in_dim %8771, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8773 = stablehlo.subtract %8772, %8770 : tensor<1xf32> loc(#loc5260)
    %8774 = stablehlo.divide %8767, %8773 : tensor<1xf32> loc(#loc5253)
    %8775 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8776 = stablehlo.broadcast_in_dim %8775, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8777 = stablehlo.subtract %8776, %8751 : tensor<1xf32> loc(#loc5260)
    %8778 = stablehlo.broadcast_in_dim %8777, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8779 = stablehlo.broadcast_in_dim %8778, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8780 = stablehlo.multiply %8779, %4500 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8781 = stablehlo.broadcast_in_dim %8751, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8782 = stablehlo.broadcast_in_dim %8781, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8783 = stablehlo.multiply %8782, %arg255 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8784 = stablehlo.add %8780, %8783 : tensor<1x2048x8192xf32> loc(#loc5263)
    %8785 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8786 = stablehlo.broadcast_in_dim %8785, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8787 = stablehlo.subtract %8786, %8774 : tensor<1xf32> loc(#loc5260)
    %8788 = stablehlo.multiply %4500, %4500 : tensor<1x2048x8192xf32> loc(#loc5264)
    %8789 = stablehlo.broadcast_in_dim %8787, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8790 = stablehlo.broadcast_in_dim %8789, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8791 = stablehlo.multiply %8790, %8788 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8792 = stablehlo.broadcast_in_dim %8774, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8793 = stablehlo.broadcast_in_dim %8792, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %8794 = stablehlo.multiply %8793, %arg431 : tensor<1x2048x8192xf32> loc(#loc5262)
    %8795 = stablehlo.add %8791, %8794 : tensor<1x2048x8192xf32> loc(#loc5263)
    %8796 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8797 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8798 = stablehlo.broadcast_in_dim %8797, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8799 = stablehlo.add %8796, %8798 : tensor<1xf32> loc(#loc5259)
    %8800 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8801 = stablehlo.broadcast_in_dim %8800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8802 = stablehlo.subtract %8799, %8801 : tensor<1xf32> loc(#loc5260)
    %8803 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8804 = stablehlo.broadcast_in_dim %8803, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8805 = stablehlo.power %8804, %8802 : tensor<1xf32> loc(#loc5261)
    %8806 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8807 = stablehlo.broadcast_in_dim %8806, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8808 = stablehlo.subtract %8807, %8805 : tensor<1xf32> loc(#loc5260)
    %8809 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8810 = stablehlo.broadcast_in_dim %8809, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8811 = stablehlo.multiply %8810, %8808 : tensor<1xf32> loc(#loc5236)
    %8812 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8813 = stablehlo.broadcast_in_dim %8812, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8814 = stablehlo.power %8813, %8799 : tensor<1xf32> loc(#loc5261)
    %8815 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8816 = stablehlo.broadcast_in_dim %8815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8817 = stablehlo.subtract %8816, %8814 : tensor<1xf32> loc(#loc5260)
    %8818 = stablehlo.divide %8811, %8817 : tensor<1xf32> loc(#loc5253)
    %8819 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8820 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8821 = stablehlo.broadcast_in_dim %8820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8822 = stablehlo.add %8819, %8821 : tensor<1xf32> loc(#loc5259)
    %8823 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8824 = stablehlo.broadcast_in_dim %8823, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8825 = stablehlo.subtract %8822, %8824 : tensor<1xf32> loc(#loc5260)
    %8826 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8827 = stablehlo.broadcast_in_dim %8826, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8828 = stablehlo.power %8827, %8825 : tensor<1xf32> loc(#loc5261)
    %8829 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8830 = stablehlo.broadcast_in_dim %8829, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8831 = stablehlo.subtract %8830, %8828 : tensor<1xf32> loc(#loc5260)
    %8832 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8833 = stablehlo.broadcast_in_dim %8832, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8834 = stablehlo.multiply %8833, %8831 : tensor<1xf32> loc(#loc5236)
    %8835 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8836 = stablehlo.broadcast_in_dim %8835, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8837 = stablehlo.power %8836, %8822 : tensor<1xf32> loc(#loc5261)
    %8838 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8839 = stablehlo.broadcast_in_dim %8838, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8840 = stablehlo.subtract %8839, %8837 : tensor<1xf32> loc(#loc5260)
    %8841 = stablehlo.divide %8834, %8840 : tensor<1xf32> loc(#loc5253)
    %8842 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8843 = stablehlo.broadcast_in_dim %8842, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8844 = stablehlo.subtract %8843, %8818 : tensor<1xf32> loc(#loc5260)
    %8845 = stablehlo.broadcast_in_dim %8844, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8846 = stablehlo.broadcast_in_dim %8845, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8847 = stablehlo.multiply %8846, %4504 : tensor<1x2048xf32> loc(#loc5262)
    %8848 = stablehlo.broadcast_in_dim %8818, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8849 = stablehlo.broadcast_in_dim %8848, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8850 = stablehlo.multiply %8849, %arg256 : tensor<1x2048xf32> loc(#loc5262)
    %8851 = stablehlo.add %8847, %8850 : tensor<1x2048xf32> loc(#loc5263)
    %8852 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8853 = stablehlo.broadcast_in_dim %8852, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8854 = stablehlo.subtract %8853, %8841 : tensor<1xf32> loc(#loc5260)
    %8855 = stablehlo.multiply %4504, %4504 : tensor<1x2048xf32> loc(#loc5264)
    %8856 = stablehlo.broadcast_in_dim %8854, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8857 = stablehlo.broadcast_in_dim %8856, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8858 = stablehlo.multiply %8857, %8855 : tensor<1x2048xf32> loc(#loc5262)
    %8859 = stablehlo.broadcast_in_dim %8841, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8860 = stablehlo.broadcast_in_dim %8859, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8861 = stablehlo.multiply %8860, %arg432 : tensor<1x2048xf32> loc(#loc5262)
    %8862 = stablehlo.add %8858, %8861 : tensor<1x2048xf32> loc(#loc5263)
    %8863 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8864 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8865 = stablehlo.broadcast_in_dim %8864, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8866 = stablehlo.add %8863, %8865 : tensor<1xf32> loc(#loc5259)
    %8867 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8868 = stablehlo.broadcast_in_dim %8867, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8869 = stablehlo.subtract %8866, %8868 : tensor<1xf32> loc(#loc5260)
    %8870 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8871 = stablehlo.broadcast_in_dim %8870, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8872 = stablehlo.power %8871, %8869 : tensor<1xf32> loc(#loc5261)
    %8873 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8874 = stablehlo.broadcast_in_dim %8873, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8875 = stablehlo.subtract %8874, %8872 : tensor<1xf32> loc(#loc5260)
    %8876 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8877 = stablehlo.broadcast_in_dim %8876, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8878 = stablehlo.multiply %8877, %8875 : tensor<1xf32> loc(#loc5236)
    %8879 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8880 = stablehlo.broadcast_in_dim %8879, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8881 = stablehlo.power %8880, %8866 : tensor<1xf32> loc(#loc5261)
    %8882 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8883 = stablehlo.broadcast_in_dim %8882, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8884 = stablehlo.subtract %8883, %8881 : tensor<1xf32> loc(#loc5260)
    %8885 = stablehlo.divide %8878, %8884 : tensor<1xf32> loc(#loc5253)
    %8886 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8887 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8888 = stablehlo.broadcast_in_dim %8887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8889 = stablehlo.add %8886, %8888 : tensor<1xf32> loc(#loc5259)
    %8890 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8891 = stablehlo.broadcast_in_dim %8890, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8892 = stablehlo.subtract %8889, %8891 : tensor<1xf32> loc(#loc5260)
    %8893 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8894 = stablehlo.broadcast_in_dim %8893, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8895 = stablehlo.power %8894, %8892 : tensor<1xf32> loc(#loc5261)
    %8896 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8897 = stablehlo.broadcast_in_dim %8896, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8898 = stablehlo.subtract %8897, %8895 : tensor<1xf32> loc(#loc5260)
    %8899 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8900 = stablehlo.broadcast_in_dim %8899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8901 = stablehlo.multiply %8900, %8898 : tensor<1xf32> loc(#loc5236)
    %8902 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8903 = stablehlo.broadcast_in_dim %8902, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8904 = stablehlo.power %8903, %8889 : tensor<1xf32> loc(#loc5261)
    %8905 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8906 = stablehlo.broadcast_in_dim %8905, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8907 = stablehlo.subtract %8906, %8904 : tensor<1xf32> loc(#loc5260)
    %8908 = stablehlo.divide %8901, %8907 : tensor<1xf32> loc(#loc5253)
    %8909 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8910 = stablehlo.broadcast_in_dim %8909, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8911 = stablehlo.subtract %8910, %8885 : tensor<1xf32> loc(#loc5260)
    %8912 = stablehlo.broadcast_in_dim %8911, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8913 = stablehlo.broadcast_in_dim %8912, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8914 = stablehlo.multiply %8913, %4508 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8915 = stablehlo.broadcast_in_dim %8885, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8916 = stablehlo.broadcast_in_dim %8915, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8917 = stablehlo.multiply %8916, %arg257 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8918 = stablehlo.add %8914, %8917 : tensor<1x8192x2048xf32> loc(#loc5263)
    %8919 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8920 = stablehlo.broadcast_in_dim %8919, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8921 = stablehlo.subtract %8920, %8908 : tensor<1xf32> loc(#loc5260)
    %8922 = stablehlo.multiply %4508, %4508 : tensor<1x8192x2048xf32> loc(#loc5264)
    %8923 = stablehlo.broadcast_in_dim %8921, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8924 = stablehlo.broadcast_in_dim %8923, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8925 = stablehlo.multiply %8924, %8922 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8926 = stablehlo.broadcast_in_dim %8908, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %8927 = stablehlo.broadcast_in_dim %8926, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %8928 = stablehlo.multiply %8927, %arg433 : tensor<1x8192x2048xf32> loc(#loc5262)
    %8929 = stablehlo.add %8925, %8928 : tensor<1x8192x2048xf32> loc(#loc5263)
    %8930 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8931 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8932 = stablehlo.broadcast_in_dim %8931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8933 = stablehlo.add %8930, %8932 : tensor<1xf32> loc(#loc5259)
    %8934 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8935 = stablehlo.broadcast_in_dim %8934, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8936 = stablehlo.subtract %8933, %8935 : tensor<1xf32> loc(#loc5260)
    %8937 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8938 = stablehlo.broadcast_in_dim %8937, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8939 = stablehlo.power %8938, %8936 : tensor<1xf32> loc(#loc5261)
    %8940 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8941 = stablehlo.broadcast_in_dim %8940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8942 = stablehlo.subtract %8941, %8939 : tensor<1xf32> loc(#loc5260)
    %8943 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8944 = stablehlo.broadcast_in_dim %8943, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8945 = stablehlo.multiply %8944, %8942 : tensor<1xf32> loc(#loc5236)
    %8946 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %8947 = stablehlo.broadcast_in_dim %8946, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8948 = stablehlo.power %8947, %8933 : tensor<1xf32> loc(#loc5261)
    %8949 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8950 = stablehlo.broadcast_in_dim %8949, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8951 = stablehlo.subtract %8950, %8948 : tensor<1xf32> loc(#loc5260)
    %8952 = stablehlo.divide %8945, %8951 : tensor<1xf32> loc(#loc5253)
    %8953 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8954 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8955 = stablehlo.broadcast_in_dim %8954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %8956 = stablehlo.add %8953, %8955 : tensor<1xf32> loc(#loc5259)
    %8957 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8958 = stablehlo.broadcast_in_dim %8957, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8959 = stablehlo.subtract %8956, %8958 : tensor<1xf32> loc(#loc5260)
    %8960 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8961 = stablehlo.broadcast_in_dim %8960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8962 = stablehlo.power %8961, %8959 : tensor<1xf32> loc(#loc5261)
    %8963 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8964 = stablehlo.broadcast_in_dim %8963, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8965 = stablehlo.subtract %8964, %8962 : tensor<1xf32> loc(#loc5260)
    %8966 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8967 = stablehlo.broadcast_in_dim %8966, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %8968 = stablehlo.multiply %8967, %8965 : tensor<1xf32> loc(#loc5236)
    %8969 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %8970 = stablehlo.broadcast_in_dim %8969, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %8971 = stablehlo.power %8970, %8956 : tensor<1xf32> loc(#loc5261)
    %8972 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8973 = stablehlo.broadcast_in_dim %8972, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8974 = stablehlo.subtract %8973, %8971 : tensor<1xf32> loc(#loc5260)
    %8975 = stablehlo.divide %8968, %8974 : tensor<1xf32> loc(#loc5253)
    %8976 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8977 = stablehlo.broadcast_in_dim %8976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8978 = stablehlo.subtract %8977, %8952 : tensor<1xf32> loc(#loc5260)
    %8979 = stablehlo.broadcast_in_dim %8978, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8980 = stablehlo.broadcast_in_dim %8979, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8981 = stablehlo.multiply %8980, %4512 : tensor<1x2048xf32> loc(#loc5262)
    %8982 = stablehlo.broadcast_in_dim %8952, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8983 = stablehlo.broadcast_in_dim %8982, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8984 = stablehlo.multiply %8983, %arg258 : tensor<1x2048xf32> loc(#loc5262)
    %8985 = stablehlo.add %8981, %8984 : tensor<1x2048xf32> loc(#loc5263)
    %8986 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8987 = stablehlo.broadcast_in_dim %8986, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %8988 = stablehlo.subtract %8987, %8975 : tensor<1xf32> loc(#loc5260)
    %8989 = stablehlo.multiply %4512, %4512 : tensor<1x2048xf32> loc(#loc5264)
    %8990 = stablehlo.broadcast_in_dim %8988, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8991 = stablehlo.broadcast_in_dim %8990, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8992 = stablehlo.multiply %8991, %8989 : tensor<1x2048xf32> loc(#loc5262)
    %8993 = stablehlo.broadcast_in_dim %8975, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %8994 = stablehlo.broadcast_in_dim %8993, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %8995 = stablehlo.multiply %8994, %arg434 : tensor<1x2048xf32> loc(#loc5262)
    %8996 = stablehlo.add %8992, %8995 : tensor<1x2048xf32> loc(#loc5263)
    %8997 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %8998 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %8999 = stablehlo.broadcast_in_dim %8998, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9000 = stablehlo.add %8997, %8999 : tensor<1xf32> loc(#loc5259)
    %9001 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9002 = stablehlo.broadcast_in_dim %9001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9003 = stablehlo.subtract %9000, %9002 : tensor<1xf32> loc(#loc5260)
    %9004 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9005 = stablehlo.broadcast_in_dim %9004, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9006 = stablehlo.power %9005, %9003 : tensor<1xf32> loc(#loc5261)
    %9007 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9008 = stablehlo.broadcast_in_dim %9007, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9009 = stablehlo.subtract %9008, %9006 : tensor<1xf32> loc(#loc5260)
    %9010 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9011 = stablehlo.broadcast_in_dim %9010, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9012 = stablehlo.multiply %9011, %9009 : tensor<1xf32> loc(#loc5236)
    %9013 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9014 = stablehlo.broadcast_in_dim %9013, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9015 = stablehlo.power %9014, %9000 : tensor<1xf32> loc(#loc5261)
    %9016 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9017 = stablehlo.broadcast_in_dim %9016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9018 = stablehlo.subtract %9017, %9015 : tensor<1xf32> loc(#loc5260)
    %9019 = stablehlo.divide %9012, %9018 : tensor<1xf32> loc(#loc5253)
    %9020 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9021 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9022 = stablehlo.broadcast_in_dim %9021, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9023 = stablehlo.add %9020, %9022 : tensor<1xf32> loc(#loc5259)
    %9024 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9025 = stablehlo.broadcast_in_dim %9024, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9026 = stablehlo.subtract %9023, %9025 : tensor<1xf32> loc(#loc5260)
    %9027 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9028 = stablehlo.broadcast_in_dim %9027, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9029 = stablehlo.power %9028, %9026 : tensor<1xf32> loc(#loc5261)
    %9030 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9031 = stablehlo.broadcast_in_dim %9030, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9032 = stablehlo.subtract %9031, %9029 : tensor<1xf32> loc(#loc5260)
    %9033 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9034 = stablehlo.broadcast_in_dim %9033, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9035 = stablehlo.multiply %9034, %9032 : tensor<1xf32> loc(#loc5236)
    %9036 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9037 = stablehlo.broadcast_in_dim %9036, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9038 = stablehlo.power %9037, %9023 : tensor<1xf32> loc(#loc5261)
    %9039 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9040 = stablehlo.broadcast_in_dim %9039, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9041 = stablehlo.subtract %9040, %9038 : tensor<1xf32> loc(#loc5260)
    %9042 = stablehlo.divide %9035, %9041 : tensor<1xf32> loc(#loc5253)
    %9043 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9044 = stablehlo.broadcast_in_dim %9043, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9045 = stablehlo.subtract %9044, %9019 : tensor<1xf32> loc(#loc5260)
    %9046 = stablehlo.broadcast_in_dim %9045, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9047 = stablehlo.broadcast_in_dim %9046, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9048 = stablehlo.multiply %9047, %4516 : tensor<1x2048xf32> loc(#loc5262)
    %9049 = stablehlo.broadcast_in_dim %9019, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9050 = stablehlo.broadcast_in_dim %9049, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9051 = stablehlo.multiply %9050, %arg259 : tensor<1x2048xf32> loc(#loc5262)
    %9052 = stablehlo.add %9048, %9051 : tensor<1x2048xf32> loc(#loc5263)
    %9053 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9054 = stablehlo.broadcast_in_dim %9053, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9055 = stablehlo.subtract %9054, %9042 : tensor<1xf32> loc(#loc5260)
    %9056 = stablehlo.multiply %4516, %4516 : tensor<1x2048xf32> loc(#loc5264)
    %9057 = stablehlo.broadcast_in_dim %9055, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9058 = stablehlo.broadcast_in_dim %9057, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9059 = stablehlo.multiply %9058, %9056 : tensor<1x2048xf32> loc(#loc5262)
    %9060 = stablehlo.broadcast_in_dim %9042, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9061 = stablehlo.broadcast_in_dim %9060, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9062 = stablehlo.multiply %9061, %arg435 : tensor<1x2048xf32> loc(#loc5262)
    %9063 = stablehlo.add %9059, %9062 : tensor<1x2048xf32> loc(#loc5263)
    %9064 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9065 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9066 = stablehlo.broadcast_in_dim %9065, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9067 = stablehlo.add %9064, %9066 : tensor<1xf32> loc(#loc5259)
    %9068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9069 = stablehlo.broadcast_in_dim %9068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9070 = stablehlo.subtract %9067, %9069 : tensor<1xf32> loc(#loc5260)
    %9071 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9072 = stablehlo.broadcast_in_dim %9071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9073 = stablehlo.power %9072, %9070 : tensor<1xf32> loc(#loc5261)
    %9074 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9075 = stablehlo.broadcast_in_dim %9074, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9076 = stablehlo.subtract %9075, %9073 : tensor<1xf32> loc(#loc5260)
    %9077 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9078 = stablehlo.broadcast_in_dim %9077, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9079 = stablehlo.multiply %9078, %9076 : tensor<1xf32> loc(#loc5236)
    %9080 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9081 = stablehlo.broadcast_in_dim %9080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9082 = stablehlo.power %9081, %9067 : tensor<1xf32> loc(#loc5261)
    %9083 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9084 = stablehlo.broadcast_in_dim %9083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9085 = stablehlo.subtract %9084, %9082 : tensor<1xf32> loc(#loc5260)
    %9086 = stablehlo.divide %9079, %9085 : tensor<1xf32> loc(#loc5253)
    %9087 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9089 = stablehlo.broadcast_in_dim %9088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9090 = stablehlo.add %9087, %9089 : tensor<1xf32> loc(#loc5259)
    %9091 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9092 = stablehlo.broadcast_in_dim %9091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9093 = stablehlo.subtract %9090, %9092 : tensor<1xf32> loc(#loc5260)
    %9094 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9095 = stablehlo.broadcast_in_dim %9094, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9096 = stablehlo.power %9095, %9093 : tensor<1xf32> loc(#loc5261)
    %9097 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9098 = stablehlo.broadcast_in_dim %9097, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9099 = stablehlo.subtract %9098, %9096 : tensor<1xf32> loc(#loc5260)
    %9100 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9101 = stablehlo.broadcast_in_dim %9100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9102 = stablehlo.multiply %9101, %9099 : tensor<1xf32> loc(#loc5236)
    %9103 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9104 = stablehlo.broadcast_in_dim %9103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9105 = stablehlo.power %9104, %9090 : tensor<1xf32> loc(#loc5261)
    %9106 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9107 = stablehlo.broadcast_in_dim %9106, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9108 = stablehlo.subtract %9107, %9105 : tensor<1xf32> loc(#loc5260)
    %9109 = stablehlo.divide %9102, %9108 : tensor<1xf32> loc(#loc5253)
    %9110 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9111 = stablehlo.broadcast_in_dim %9110, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9112 = stablehlo.subtract %9111, %9086 : tensor<1xf32> loc(#loc5260)
    %9113 = stablehlo.broadcast_in_dim %9112, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9114 = stablehlo.broadcast_in_dim %9113, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9115 = stablehlo.multiply %9114, %4520 : tensor<1x2048xf32> loc(#loc5262)
    %9116 = stablehlo.broadcast_in_dim %9086, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9117 = stablehlo.broadcast_in_dim %9116, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9118 = stablehlo.multiply %9117, %arg260 : tensor<1x2048xf32> loc(#loc5262)
    %9119 = stablehlo.add %9115, %9118 : tensor<1x2048xf32> loc(#loc5263)
    %9120 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9121 = stablehlo.broadcast_in_dim %9120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9122 = stablehlo.subtract %9121, %9109 : tensor<1xf32> loc(#loc5260)
    %9123 = stablehlo.multiply %4520, %4520 : tensor<1x2048xf32> loc(#loc5264)
    %9124 = stablehlo.broadcast_in_dim %9122, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9125 = stablehlo.broadcast_in_dim %9124, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9126 = stablehlo.multiply %9125, %9123 : tensor<1x2048xf32> loc(#loc5262)
    %9127 = stablehlo.broadcast_in_dim %9109, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9128 = stablehlo.broadcast_in_dim %9127, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9129 = stablehlo.multiply %9128, %arg436 : tensor<1x2048xf32> loc(#loc5262)
    %9130 = stablehlo.add %9126, %9129 : tensor<1x2048xf32> loc(#loc5263)
    %9131 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9132 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9133 = stablehlo.broadcast_in_dim %9132, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9134 = stablehlo.add %9131, %9133 : tensor<1xf32> loc(#loc5259)
    %9135 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9136 = stablehlo.broadcast_in_dim %9135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9137 = stablehlo.subtract %9134, %9136 : tensor<1xf32> loc(#loc5260)
    %9138 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9139 = stablehlo.broadcast_in_dim %9138, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9140 = stablehlo.power %9139, %9137 : tensor<1xf32> loc(#loc5261)
    %9141 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9142 = stablehlo.broadcast_in_dim %9141, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9143 = stablehlo.subtract %9142, %9140 : tensor<1xf32> loc(#loc5260)
    %9144 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9145 = stablehlo.broadcast_in_dim %9144, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9146 = stablehlo.multiply %9145, %9143 : tensor<1xf32> loc(#loc5236)
    %9147 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9148 = stablehlo.broadcast_in_dim %9147, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9149 = stablehlo.power %9148, %9134 : tensor<1xf32> loc(#loc5261)
    %9150 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9151 = stablehlo.broadcast_in_dim %9150, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9152 = stablehlo.subtract %9151, %9149 : tensor<1xf32> loc(#loc5260)
    %9153 = stablehlo.divide %9146, %9152 : tensor<1xf32> loc(#loc5253)
    %9154 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9155 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9156 = stablehlo.broadcast_in_dim %9155, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9157 = stablehlo.add %9154, %9156 : tensor<1xf32> loc(#loc5259)
    %9158 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9159 = stablehlo.broadcast_in_dim %9158, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9160 = stablehlo.subtract %9157, %9159 : tensor<1xf32> loc(#loc5260)
    %9161 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9162 = stablehlo.broadcast_in_dim %9161, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9163 = stablehlo.power %9162, %9160 : tensor<1xf32> loc(#loc5261)
    %9164 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9165 = stablehlo.broadcast_in_dim %9164, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9166 = stablehlo.subtract %9165, %9163 : tensor<1xf32> loc(#loc5260)
    %9167 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9168 = stablehlo.broadcast_in_dim %9167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9169 = stablehlo.multiply %9168, %9166 : tensor<1xf32> loc(#loc5236)
    %9170 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9171 = stablehlo.broadcast_in_dim %9170, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9172 = stablehlo.power %9171, %9157 : tensor<1xf32> loc(#loc5261)
    %9173 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9174 = stablehlo.broadcast_in_dim %9173, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9175 = stablehlo.subtract %9174, %9172 : tensor<1xf32> loc(#loc5260)
    %9176 = stablehlo.divide %9169, %9175 : tensor<1xf32> loc(#loc5253)
    %9177 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9178 = stablehlo.broadcast_in_dim %9177, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9179 = stablehlo.subtract %9178, %9153 : tensor<1xf32> loc(#loc5260)
    %9180 = stablehlo.broadcast_in_dim %9179, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9181 = stablehlo.broadcast_in_dim %9180, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9182 = stablehlo.multiply %9181, %4524 : tensor<1x2048xf32> loc(#loc5262)
    %9183 = stablehlo.broadcast_in_dim %9153, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9184 = stablehlo.broadcast_in_dim %9183, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9185 = stablehlo.multiply %9184, %arg261 : tensor<1x2048xf32> loc(#loc5262)
    %9186 = stablehlo.add %9182, %9185 : tensor<1x2048xf32> loc(#loc5263)
    %9187 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9188 = stablehlo.broadcast_in_dim %9187, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9189 = stablehlo.subtract %9188, %9176 : tensor<1xf32> loc(#loc5260)
    %9190 = stablehlo.multiply %4524, %4524 : tensor<1x2048xf32> loc(#loc5264)
    %9191 = stablehlo.broadcast_in_dim %9189, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9192 = stablehlo.broadcast_in_dim %9191, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9193 = stablehlo.multiply %9192, %9190 : tensor<1x2048xf32> loc(#loc5262)
    %9194 = stablehlo.broadcast_in_dim %9176, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9195 = stablehlo.broadcast_in_dim %9194, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9196 = stablehlo.multiply %9195, %arg437 : tensor<1x2048xf32> loc(#loc5262)
    %9197 = stablehlo.add %9193, %9196 : tensor<1x2048xf32> loc(#loc5263)
    %9198 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9199 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9200 = stablehlo.broadcast_in_dim %9199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9201 = stablehlo.add %9198, %9200 : tensor<1xf32> loc(#loc5259)
    %9202 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9203 = stablehlo.broadcast_in_dim %9202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9204 = stablehlo.subtract %9201, %9203 : tensor<1xf32> loc(#loc5260)
    %9205 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9206 = stablehlo.broadcast_in_dim %9205, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9207 = stablehlo.power %9206, %9204 : tensor<1xf32> loc(#loc5261)
    %9208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9209 = stablehlo.broadcast_in_dim %9208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9210 = stablehlo.subtract %9209, %9207 : tensor<1xf32> loc(#loc5260)
    %9211 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9212 = stablehlo.broadcast_in_dim %9211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9213 = stablehlo.multiply %9212, %9210 : tensor<1xf32> loc(#loc5236)
    %9214 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9215 = stablehlo.broadcast_in_dim %9214, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9216 = stablehlo.power %9215, %9201 : tensor<1xf32> loc(#loc5261)
    %9217 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9218 = stablehlo.broadcast_in_dim %9217, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9219 = stablehlo.subtract %9218, %9216 : tensor<1xf32> loc(#loc5260)
    %9220 = stablehlo.divide %9213, %9219 : tensor<1xf32> loc(#loc5253)
    %9221 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9222 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9223 = stablehlo.broadcast_in_dim %9222, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9224 = stablehlo.add %9221, %9223 : tensor<1xf32> loc(#loc5259)
    %9225 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9226 = stablehlo.broadcast_in_dim %9225, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9227 = stablehlo.subtract %9224, %9226 : tensor<1xf32> loc(#loc5260)
    %9228 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9229 = stablehlo.broadcast_in_dim %9228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9230 = stablehlo.power %9229, %9227 : tensor<1xf32> loc(#loc5261)
    %9231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9232 = stablehlo.broadcast_in_dim %9231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9233 = stablehlo.subtract %9232, %9230 : tensor<1xf32> loc(#loc5260)
    %9234 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9235 = stablehlo.broadcast_in_dim %9234, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9236 = stablehlo.multiply %9235, %9233 : tensor<1xf32> loc(#loc5236)
    %9237 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9238 = stablehlo.broadcast_in_dim %9237, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9239 = stablehlo.power %9238, %9224 : tensor<1xf32> loc(#loc5261)
    %9240 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9241 = stablehlo.broadcast_in_dim %9240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9242 = stablehlo.subtract %9241, %9239 : tensor<1xf32> loc(#loc5260)
    %9243 = stablehlo.divide %9236, %9242 : tensor<1xf32> loc(#loc5253)
    %9244 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9245 = stablehlo.broadcast_in_dim %9244, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9246 = stablehlo.subtract %9245, %9220 : tensor<1xf32> loc(#loc5260)
    %9247 = stablehlo.broadcast_in_dim %9246, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9248 = stablehlo.broadcast_in_dim %9247, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9249 = stablehlo.multiply %9248, %4528 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9250 = stablehlo.broadcast_in_dim %9220, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9251 = stablehlo.broadcast_in_dim %9250, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9252 = stablehlo.multiply %9251, %arg262 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9253 = stablehlo.add %9249, %9252 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %9254 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9255 = stablehlo.broadcast_in_dim %9254, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9256 = stablehlo.subtract %9255, %9243 : tensor<1xf32> loc(#loc5260)
    %9257 = stablehlo.multiply %4528, %4528 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %9258 = stablehlo.broadcast_in_dim %9256, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9259 = stablehlo.broadcast_in_dim %9258, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9260 = stablehlo.multiply %9259, %9257 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9261 = stablehlo.broadcast_in_dim %9243, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9262 = stablehlo.broadcast_in_dim %9261, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9263 = stablehlo.multiply %9262, %arg438 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9264 = stablehlo.add %9260, %9263 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %9265 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9266 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9267 = stablehlo.broadcast_in_dim %9266, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9268 = stablehlo.add %9265, %9267 : tensor<1xf32> loc(#loc5259)
    %9269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9270 = stablehlo.broadcast_in_dim %9269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9271 = stablehlo.subtract %9268, %9270 : tensor<1xf32> loc(#loc5260)
    %9272 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9273 = stablehlo.broadcast_in_dim %9272, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9274 = stablehlo.power %9273, %9271 : tensor<1xf32> loc(#loc5261)
    %9275 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9276 = stablehlo.broadcast_in_dim %9275, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9277 = stablehlo.subtract %9276, %9274 : tensor<1xf32> loc(#loc5260)
    %9278 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9279 = stablehlo.broadcast_in_dim %9278, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9280 = stablehlo.multiply %9279, %9277 : tensor<1xf32> loc(#loc5236)
    %9281 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9282 = stablehlo.broadcast_in_dim %9281, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9283 = stablehlo.power %9282, %9268 : tensor<1xf32> loc(#loc5261)
    %9284 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9285 = stablehlo.broadcast_in_dim %9284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9286 = stablehlo.subtract %9285, %9283 : tensor<1xf32> loc(#loc5260)
    %9287 = stablehlo.divide %9280, %9286 : tensor<1xf32> loc(#loc5253)
    %9288 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9289 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9290 = stablehlo.broadcast_in_dim %9289, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9291 = stablehlo.add %9288, %9290 : tensor<1xf32> loc(#loc5259)
    %9292 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9293 = stablehlo.broadcast_in_dim %9292, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9294 = stablehlo.subtract %9291, %9293 : tensor<1xf32> loc(#loc5260)
    %9295 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9296 = stablehlo.broadcast_in_dim %9295, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9297 = stablehlo.power %9296, %9294 : tensor<1xf32> loc(#loc5261)
    %9298 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9299 = stablehlo.broadcast_in_dim %9298, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9300 = stablehlo.subtract %9299, %9297 : tensor<1xf32> loc(#loc5260)
    %9301 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9302 = stablehlo.broadcast_in_dim %9301, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9303 = stablehlo.multiply %9302, %9300 : tensor<1xf32> loc(#loc5236)
    %9304 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9305 = stablehlo.broadcast_in_dim %9304, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9306 = stablehlo.power %9305, %9291 : tensor<1xf32> loc(#loc5261)
    %9307 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9308 = stablehlo.broadcast_in_dim %9307, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9309 = stablehlo.subtract %9308, %9306 : tensor<1xf32> loc(#loc5260)
    %9310 = stablehlo.divide %9303, %9309 : tensor<1xf32> loc(#loc5253)
    %9311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9312 = stablehlo.broadcast_in_dim %9311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9313 = stablehlo.subtract %9312, %9287 : tensor<1xf32> loc(#loc5260)
    %9314 = stablehlo.broadcast_in_dim %9313, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9315 = stablehlo.broadcast_in_dim %9314, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %9316 = stablehlo.multiply %9315, %4532 : tensor<1x64xf32> loc(#loc5262)
    %9317 = stablehlo.broadcast_in_dim %9287, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9318 = stablehlo.broadcast_in_dim %9317, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %9319 = stablehlo.multiply %9318, %arg263 : tensor<1x64xf32> loc(#loc5262)
    %9320 = stablehlo.add %9316, %9319 : tensor<1x64xf32> loc(#loc5263)
    %9321 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9322 = stablehlo.broadcast_in_dim %9321, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9323 = stablehlo.subtract %9322, %9310 : tensor<1xf32> loc(#loc5260)
    %9324 = stablehlo.multiply %4532, %4532 : tensor<1x64xf32> loc(#loc5264)
    %9325 = stablehlo.broadcast_in_dim %9323, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9326 = stablehlo.broadcast_in_dim %9325, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %9327 = stablehlo.multiply %9326, %9324 : tensor<1x64xf32> loc(#loc5262)
    %9328 = stablehlo.broadcast_in_dim %9310, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9329 = stablehlo.broadcast_in_dim %9328, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %9330 = stablehlo.multiply %9329, %arg439 : tensor<1x64xf32> loc(#loc5262)
    %9331 = stablehlo.add %9327, %9330 : tensor<1x64xf32> loc(#loc5263)
    %9332 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9333 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9334 = stablehlo.broadcast_in_dim %9333, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9335 = stablehlo.add %9332, %9334 : tensor<1xf32> loc(#loc5259)
    %9336 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9337 = stablehlo.broadcast_in_dim %9336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9338 = stablehlo.subtract %9335, %9337 : tensor<1xf32> loc(#loc5260)
    %9339 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9340 = stablehlo.broadcast_in_dim %9339, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9341 = stablehlo.power %9340, %9338 : tensor<1xf32> loc(#loc5261)
    %9342 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9343 = stablehlo.broadcast_in_dim %9342, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9344 = stablehlo.subtract %9343, %9341 : tensor<1xf32> loc(#loc5260)
    %9345 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9346 = stablehlo.broadcast_in_dim %9345, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9347 = stablehlo.multiply %9346, %9344 : tensor<1xf32> loc(#loc5236)
    %9348 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9349 = stablehlo.broadcast_in_dim %9348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9350 = stablehlo.power %9349, %9335 : tensor<1xf32> loc(#loc5261)
    %9351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9352 = stablehlo.broadcast_in_dim %9351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9353 = stablehlo.subtract %9352, %9350 : tensor<1xf32> loc(#loc5260)
    %9354 = stablehlo.divide %9347, %9353 : tensor<1xf32> loc(#loc5253)
    %9355 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9356 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9357 = stablehlo.broadcast_in_dim %9356, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9358 = stablehlo.add %9355, %9357 : tensor<1xf32> loc(#loc5259)
    %9359 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9360 = stablehlo.broadcast_in_dim %9359, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9361 = stablehlo.subtract %9358, %9360 : tensor<1xf32> loc(#loc5260)
    %9362 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9363 = stablehlo.broadcast_in_dim %9362, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9364 = stablehlo.power %9363, %9361 : tensor<1xf32> loc(#loc5261)
    %9365 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9366 = stablehlo.broadcast_in_dim %9365, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9367 = stablehlo.subtract %9366, %9364 : tensor<1xf32> loc(#loc5260)
    %9368 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9369 = stablehlo.broadcast_in_dim %9368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9370 = stablehlo.multiply %9369, %9367 : tensor<1xf32> loc(#loc5236)
    %9371 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9372 = stablehlo.broadcast_in_dim %9371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9373 = stablehlo.power %9372, %9358 : tensor<1xf32> loc(#loc5261)
    %9374 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9375 = stablehlo.broadcast_in_dim %9374, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9376 = stablehlo.subtract %9375, %9373 : tensor<1xf32> loc(#loc5260)
    %9377 = stablehlo.divide %9370, %9376 : tensor<1xf32> loc(#loc5253)
    %9378 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9379 = stablehlo.broadcast_in_dim %9378, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9380 = stablehlo.subtract %9379, %9354 : tensor<1xf32> loc(#loc5260)
    %9381 = stablehlo.broadcast_in_dim %9380, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %9382 = stablehlo.broadcast_in_dim %9381, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %9383 = stablehlo.multiply %9382, %4536 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %9384 = stablehlo.broadcast_in_dim %9354, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %9385 = stablehlo.broadcast_in_dim %9384, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %9386 = stablehlo.multiply %9385, %arg264 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %9387 = stablehlo.add %9383, %9386 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %9388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9389 = stablehlo.broadcast_in_dim %9388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9390 = stablehlo.subtract %9389, %9377 : tensor<1xf32> loc(#loc5260)
    %9391 = stablehlo.multiply %4536, %4536 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %9392 = stablehlo.broadcast_in_dim %9390, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %9393 = stablehlo.broadcast_in_dim %9392, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %9394 = stablehlo.multiply %9393, %9391 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %9395 = stablehlo.broadcast_in_dim %9377, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %9396 = stablehlo.broadcast_in_dim %9395, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %9397 = stablehlo.multiply %9396, %arg440 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %9398 = stablehlo.add %9394, %9397 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %9399 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9400 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9401 = stablehlo.broadcast_in_dim %9400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9402 = stablehlo.add %9399, %9401 : tensor<1xf32> loc(#loc5259)
    %9403 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9404 = stablehlo.broadcast_in_dim %9403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9405 = stablehlo.subtract %9402, %9404 : tensor<1xf32> loc(#loc5260)
    %9406 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9407 = stablehlo.broadcast_in_dim %9406, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9408 = stablehlo.power %9407, %9405 : tensor<1xf32> loc(#loc5261)
    %9409 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9410 = stablehlo.broadcast_in_dim %9409, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9411 = stablehlo.subtract %9410, %9408 : tensor<1xf32> loc(#loc5260)
    %9412 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9413 = stablehlo.broadcast_in_dim %9412, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9414 = stablehlo.multiply %9413, %9411 : tensor<1xf32> loc(#loc5236)
    %9415 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9416 = stablehlo.broadcast_in_dim %9415, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9417 = stablehlo.power %9416, %9402 : tensor<1xf32> loc(#loc5261)
    %9418 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9419 = stablehlo.broadcast_in_dim %9418, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9420 = stablehlo.subtract %9419, %9417 : tensor<1xf32> loc(#loc5260)
    %9421 = stablehlo.divide %9414, %9420 : tensor<1xf32> loc(#loc5253)
    %9422 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9423 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9424 = stablehlo.broadcast_in_dim %9423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9425 = stablehlo.add %9422, %9424 : tensor<1xf32> loc(#loc5259)
    %9426 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9427 = stablehlo.broadcast_in_dim %9426, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9428 = stablehlo.subtract %9425, %9427 : tensor<1xf32> loc(#loc5260)
    %9429 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9430 = stablehlo.broadcast_in_dim %9429, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9431 = stablehlo.power %9430, %9428 : tensor<1xf32> loc(#loc5261)
    %9432 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9433 = stablehlo.broadcast_in_dim %9432, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9434 = stablehlo.subtract %9433, %9431 : tensor<1xf32> loc(#loc5260)
    %9435 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9436 = stablehlo.broadcast_in_dim %9435, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9437 = stablehlo.multiply %9436, %9434 : tensor<1xf32> loc(#loc5236)
    %9438 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9439 = stablehlo.broadcast_in_dim %9438, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9440 = stablehlo.power %9439, %9425 : tensor<1xf32> loc(#loc5261)
    %9441 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9442 = stablehlo.broadcast_in_dim %9441, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9443 = stablehlo.subtract %9442, %9440 : tensor<1xf32> loc(#loc5260)
    %9444 = stablehlo.divide %9437, %9443 : tensor<1xf32> loc(#loc5253)
    %9445 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9446 = stablehlo.broadcast_in_dim %9445, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9447 = stablehlo.subtract %9446, %9421 : tensor<1xf32> loc(#loc5260)
    %9448 = stablehlo.broadcast_in_dim %9447, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9449 = stablehlo.broadcast_in_dim %9448, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %9450 = stablehlo.multiply %9449, %4540 : tensor<1x8192xf32> loc(#loc5262)
    %9451 = stablehlo.broadcast_in_dim %9421, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9452 = stablehlo.broadcast_in_dim %9451, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %9453 = stablehlo.multiply %9452, %arg265 : tensor<1x8192xf32> loc(#loc5262)
    %9454 = stablehlo.add %9450, %9453 : tensor<1x8192xf32> loc(#loc5263)
    %9455 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9456 = stablehlo.broadcast_in_dim %9455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9457 = stablehlo.subtract %9456, %9444 : tensor<1xf32> loc(#loc5260)
    %9458 = stablehlo.multiply %4540, %4540 : tensor<1x8192xf32> loc(#loc5264)
    %9459 = stablehlo.broadcast_in_dim %9457, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9460 = stablehlo.broadcast_in_dim %9459, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %9461 = stablehlo.multiply %9460, %9458 : tensor<1x8192xf32> loc(#loc5262)
    %9462 = stablehlo.broadcast_in_dim %9444, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9463 = stablehlo.broadcast_in_dim %9462, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %9464 = stablehlo.multiply %9463, %arg441 : tensor<1x8192xf32> loc(#loc5262)
    %9465 = stablehlo.add %9461, %9464 : tensor<1x8192xf32> loc(#loc5263)
    %9466 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9467 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9468 = stablehlo.broadcast_in_dim %9467, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9469 = stablehlo.add %9466, %9468 : tensor<1xf32> loc(#loc5259)
    %9470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9471 = stablehlo.broadcast_in_dim %9470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9472 = stablehlo.subtract %9469, %9471 : tensor<1xf32> loc(#loc5260)
    %9473 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9474 = stablehlo.broadcast_in_dim %9473, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9475 = stablehlo.power %9474, %9472 : tensor<1xf32> loc(#loc5261)
    %9476 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9477 = stablehlo.broadcast_in_dim %9476, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9478 = stablehlo.subtract %9477, %9475 : tensor<1xf32> loc(#loc5260)
    %9479 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9480 = stablehlo.broadcast_in_dim %9479, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9481 = stablehlo.multiply %9480, %9478 : tensor<1xf32> loc(#loc5236)
    %9482 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9483 = stablehlo.broadcast_in_dim %9482, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9484 = stablehlo.power %9483, %9469 : tensor<1xf32> loc(#loc5261)
    %9485 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9486 = stablehlo.broadcast_in_dim %9485, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9487 = stablehlo.subtract %9486, %9484 : tensor<1xf32> loc(#loc5260)
    %9488 = stablehlo.divide %9481, %9487 : tensor<1xf32> loc(#loc5253)
    %9489 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9490 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9491 = stablehlo.broadcast_in_dim %9490, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9492 = stablehlo.add %9489, %9491 : tensor<1xf32> loc(#loc5259)
    %9493 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9494 = stablehlo.broadcast_in_dim %9493, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9495 = stablehlo.subtract %9492, %9494 : tensor<1xf32> loc(#loc5260)
    %9496 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9497 = stablehlo.broadcast_in_dim %9496, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9498 = stablehlo.power %9497, %9495 : tensor<1xf32> loc(#loc5261)
    %9499 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9500 = stablehlo.broadcast_in_dim %9499, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9501 = stablehlo.subtract %9500, %9498 : tensor<1xf32> loc(#loc5260)
    %9502 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9503 = stablehlo.broadcast_in_dim %9502, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9504 = stablehlo.multiply %9503, %9501 : tensor<1xf32> loc(#loc5236)
    %9505 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9506 = stablehlo.broadcast_in_dim %9505, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9507 = stablehlo.power %9506, %9492 : tensor<1xf32> loc(#loc5261)
    %9508 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9509 = stablehlo.broadcast_in_dim %9508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9510 = stablehlo.subtract %9509, %9507 : tensor<1xf32> loc(#loc5260)
    %9511 = stablehlo.divide %9504, %9510 : tensor<1xf32> loc(#loc5253)
    %9512 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9513 = stablehlo.broadcast_in_dim %9512, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9514 = stablehlo.subtract %9513, %9488 : tensor<1xf32> loc(#loc5260)
    %9515 = stablehlo.broadcast_in_dim %9514, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9516 = stablehlo.broadcast_in_dim %9515, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %9517 = stablehlo.multiply %9516, %4544 : tensor<1x2048x8192xf32> loc(#loc5262)
    %9518 = stablehlo.broadcast_in_dim %9488, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9519 = stablehlo.broadcast_in_dim %9518, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %9520 = stablehlo.multiply %9519, %arg266 : tensor<1x2048x8192xf32> loc(#loc5262)
    %9521 = stablehlo.add %9517, %9520 : tensor<1x2048x8192xf32> loc(#loc5263)
    %9522 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9523 = stablehlo.broadcast_in_dim %9522, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9524 = stablehlo.subtract %9523, %9511 : tensor<1xf32> loc(#loc5260)
    %9525 = stablehlo.multiply %4544, %4544 : tensor<1x2048x8192xf32> loc(#loc5264)
    %9526 = stablehlo.broadcast_in_dim %9524, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9527 = stablehlo.broadcast_in_dim %9526, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %9528 = stablehlo.multiply %9527, %9525 : tensor<1x2048x8192xf32> loc(#loc5262)
    %9529 = stablehlo.broadcast_in_dim %9511, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9530 = stablehlo.broadcast_in_dim %9529, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %9531 = stablehlo.multiply %9530, %arg442 : tensor<1x2048x8192xf32> loc(#loc5262)
    %9532 = stablehlo.add %9528, %9531 : tensor<1x2048x8192xf32> loc(#loc5263)
    %9533 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9534 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9535 = stablehlo.broadcast_in_dim %9534, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9536 = stablehlo.add %9533, %9535 : tensor<1xf32> loc(#loc5259)
    %9537 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9538 = stablehlo.broadcast_in_dim %9537, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9539 = stablehlo.subtract %9536, %9538 : tensor<1xf32> loc(#loc5260)
    %9540 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9541 = stablehlo.broadcast_in_dim %9540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9542 = stablehlo.power %9541, %9539 : tensor<1xf32> loc(#loc5261)
    %9543 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9544 = stablehlo.broadcast_in_dim %9543, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9545 = stablehlo.subtract %9544, %9542 : tensor<1xf32> loc(#loc5260)
    %9546 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9547 = stablehlo.broadcast_in_dim %9546, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9548 = stablehlo.multiply %9547, %9545 : tensor<1xf32> loc(#loc5236)
    %9549 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9550 = stablehlo.broadcast_in_dim %9549, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9551 = stablehlo.power %9550, %9536 : tensor<1xf32> loc(#loc5261)
    %9552 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9553 = stablehlo.broadcast_in_dim %9552, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9554 = stablehlo.subtract %9553, %9551 : tensor<1xf32> loc(#loc5260)
    %9555 = stablehlo.divide %9548, %9554 : tensor<1xf32> loc(#loc5253)
    %9556 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9557 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9558 = stablehlo.broadcast_in_dim %9557, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9559 = stablehlo.add %9556, %9558 : tensor<1xf32> loc(#loc5259)
    %9560 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9561 = stablehlo.broadcast_in_dim %9560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9562 = stablehlo.subtract %9559, %9561 : tensor<1xf32> loc(#loc5260)
    %9563 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9564 = stablehlo.broadcast_in_dim %9563, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9565 = stablehlo.power %9564, %9562 : tensor<1xf32> loc(#loc5261)
    %9566 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9567 = stablehlo.broadcast_in_dim %9566, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9568 = stablehlo.subtract %9567, %9565 : tensor<1xf32> loc(#loc5260)
    %9569 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9570 = stablehlo.broadcast_in_dim %9569, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9571 = stablehlo.multiply %9570, %9568 : tensor<1xf32> loc(#loc5236)
    %9572 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9573 = stablehlo.broadcast_in_dim %9572, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9574 = stablehlo.power %9573, %9559 : tensor<1xf32> loc(#loc5261)
    %9575 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9576 = stablehlo.broadcast_in_dim %9575, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9577 = stablehlo.subtract %9576, %9574 : tensor<1xf32> loc(#loc5260)
    %9578 = stablehlo.divide %9571, %9577 : tensor<1xf32> loc(#loc5253)
    %9579 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9580 = stablehlo.broadcast_in_dim %9579, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9581 = stablehlo.subtract %9580, %9555 : tensor<1xf32> loc(#loc5260)
    %9582 = stablehlo.broadcast_in_dim %9581, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9583 = stablehlo.broadcast_in_dim %9582, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9584 = stablehlo.multiply %9583, %4548 : tensor<1x2048xf32> loc(#loc5262)
    %9585 = stablehlo.broadcast_in_dim %9555, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9586 = stablehlo.broadcast_in_dim %9585, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9587 = stablehlo.multiply %9586, %arg267 : tensor<1x2048xf32> loc(#loc5262)
    %9588 = stablehlo.add %9584, %9587 : tensor<1x2048xf32> loc(#loc5263)
    %9589 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9590 = stablehlo.broadcast_in_dim %9589, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9591 = stablehlo.subtract %9590, %9578 : tensor<1xf32> loc(#loc5260)
    %9592 = stablehlo.multiply %4548, %4548 : tensor<1x2048xf32> loc(#loc5264)
    %9593 = stablehlo.broadcast_in_dim %9591, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9594 = stablehlo.broadcast_in_dim %9593, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9595 = stablehlo.multiply %9594, %9592 : tensor<1x2048xf32> loc(#loc5262)
    %9596 = stablehlo.broadcast_in_dim %9578, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9597 = stablehlo.broadcast_in_dim %9596, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9598 = stablehlo.multiply %9597, %arg443 : tensor<1x2048xf32> loc(#loc5262)
    %9599 = stablehlo.add %9595, %9598 : tensor<1x2048xf32> loc(#loc5263)
    %9600 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9601 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9602 = stablehlo.broadcast_in_dim %9601, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9603 = stablehlo.add %9600, %9602 : tensor<1xf32> loc(#loc5259)
    %9604 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9605 = stablehlo.broadcast_in_dim %9604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9606 = stablehlo.subtract %9603, %9605 : tensor<1xf32> loc(#loc5260)
    %9607 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9608 = stablehlo.broadcast_in_dim %9607, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9609 = stablehlo.power %9608, %9606 : tensor<1xf32> loc(#loc5261)
    %9610 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9611 = stablehlo.broadcast_in_dim %9610, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9612 = stablehlo.subtract %9611, %9609 : tensor<1xf32> loc(#loc5260)
    %9613 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9614 = stablehlo.broadcast_in_dim %9613, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9615 = stablehlo.multiply %9614, %9612 : tensor<1xf32> loc(#loc5236)
    %9616 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9617 = stablehlo.broadcast_in_dim %9616, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9618 = stablehlo.power %9617, %9603 : tensor<1xf32> loc(#loc5261)
    %9619 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9620 = stablehlo.broadcast_in_dim %9619, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9621 = stablehlo.subtract %9620, %9618 : tensor<1xf32> loc(#loc5260)
    %9622 = stablehlo.divide %9615, %9621 : tensor<1xf32> loc(#loc5253)
    %9623 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9624 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9625 = stablehlo.broadcast_in_dim %9624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9626 = stablehlo.add %9623, %9625 : tensor<1xf32> loc(#loc5259)
    %9627 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9628 = stablehlo.broadcast_in_dim %9627, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9629 = stablehlo.subtract %9626, %9628 : tensor<1xf32> loc(#loc5260)
    %9630 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9631 = stablehlo.broadcast_in_dim %9630, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9632 = stablehlo.power %9631, %9629 : tensor<1xf32> loc(#loc5261)
    %9633 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9634 = stablehlo.broadcast_in_dim %9633, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9635 = stablehlo.subtract %9634, %9632 : tensor<1xf32> loc(#loc5260)
    %9636 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9637 = stablehlo.broadcast_in_dim %9636, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9638 = stablehlo.multiply %9637, %9635 : tensor<1xf32> loc(#loc5236)
    %9639 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9640 = stablehlo.broadcast_in_dim %9639, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9641 = stablehlo.power %9640, %9626 : tensor<1xf32> loc(#loc5261)
    %9642 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9643 = stablehlo.broadcast_in_dim %9642, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9644 = stablehlo.subtract %9643, %9641 : tensor<1xf32> loc(#loc5260)
    %9645 = stablehlo.divide %9638, %9644 : tensor<1xf32> loc(#loc5253)
    %9646 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9647 = stablehlo.broadcast_in_dim %9646, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9648 = stablehlo.subtract %9647, %9622 : tensor<1xf32> loc(#loc5260)
    %9649 = stablehlo.broadcast_in_dim %9648, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9650 = stablehlo.broadcast_in_dim %9649, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %9651 = stablehlo.multiply %9650, %4552 : tensor<1x8192x2048xf32> loc(#loc5262)
    %9652 = stablehlo.broadcast_in_dim %9622, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9653 = stablehlo.broadcast_in_dim %9652, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %9654 = stablehlo.multiply %9653, %arg268 : tensor<1x8192x2048xf32> loc(#loc5262)
    %9655 = stablehlo.add %9651, %9654 : tensor<1x8192x2048xf32> loc(#loc5263)
    %9656 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9657 = stablehlo.broadcast_in_dim %9656, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9658 = stablehlo.subtract %9657, %9645 : tensor<1xf32> loc(#loc5260)
    %9659 = stablehlo.multiply %4552, %4552 : tensor<1x8192x2048xf32> loc(#loc5264)
    %9660 = stablehlo.broadcast_in_dim %9658, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9661 = stablehlo.broadcast_in_dim %9660, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %9662 = stablehlo.multiply %9661, %9659 : tensor<1x8192x2048xf32> loc(#loc5262)
    %9663 = stablehlo.broadcast_in_dim %9645, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %9664 = stablehlo.broadcast_in_dim %9663, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %9665 = stablehlo.multiply %9664, %arg444 : tensor<1x8192x2048xf32> loc(#loc5262)
    %9666 = stablehlo.add %9662, %9665 : tensor<1x8192x2048xf32> loc(#loc5263)
    %9667 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9668 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9669 = stablehlo.broadcast_in_dim %9668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9670 = stablehlo.add %9667, %9669 : tensor<1xf32> loc(#loc5259)
    %9671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9672 = stablehlo.broadcast_in_dim %9671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9673 = stablehlo.subtract %9670, %9672 : tensor<1xf32> loc(#loc5260)
    %9674 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9675 = stablehlo.broadcast_in_dim %9674, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9676 = stablehlo.power %9675, %9673 : tensor<1xf32> loc(#loc5261)
    %9677 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9678 = stablehlo.broadcast_in_dim %9677, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9679 = stablehlo.subtract %9678, %9676 : tensor<1xf32> loc(#loc5260)
    %9680 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9681 = stablehlo.broadcast_in_dim %9680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9682 = stablehlo.multiply %9681, %9679 : tensor<1xf32> loc(#loc5236)
    %9683 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9684 = stablehlo.broadcast_in_dim %9683, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9685 = stablehlo.power %9684, %9670 : tensor<1xf32> loc(#loc5261)
    %9686 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9687 = stablehlo.broadcast_in_dim %9686, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9688 = stablehlo.subtract %9687, %9685 : tensor<1xf32> loc(#loc5260)
    %9689 = stablehlo.divide %9682, %9688 : tensor<1xf32> loc(#loc5253)
    %9690 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9692 = stablehlo.broadcast_in_dim %9691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9693 = stablehlo.add %9690, %9692 : tensor<1xf32> loc(#loc5259)
    %9694 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9695 = stablehlo.broadcast_in_dim %9694, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9696 = stablehlo.subtract %9693, %9695 : tensor<1xf32> loc(#loc5260)
    %9697 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9698 = stablehlo.broadcast_in_dim %9697, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9699 = stablehlo.power %9698, %9696 : tensor<1xf32> loc(#loc5261)
    %9700 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9701 = stablehlo.broadcast_in_dim %9700, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9702 = stablehlo.subtract %9701, %9699 : tensor<1xf32> loc(#loc5260)
    %9703 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9704 = stablehlo.broadcast_in_dim %9703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9705 = stablehlo.multiply %9704, %9702 : tensor<1xf32> loc(#loc5236)
    %9706 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9707 = stablehlo.broadcast_in_dim %9706, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9708 = stablehlo.power %9707, %9693 : tensor<1xf32> loc(#loc5261)
    %9709 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9710 = stablehlo.broadcast_in_dim %9709, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9711 = stablehlo.subtract %9710, %9708 : tensor<1xf32> loc(#loc5260)
    %9712 = stablehlo.divide %9705, %9711 : tensor<1xf32> loc(#loc5253)
    %9713 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9714 = stablehlo.broadcast_in_dim %9713, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9715 = stablehlo.subtract %9714, %9689 : tensor<1xf32> loc(#loc5260)
    %9716 = stablehlo.broadcast_in_dim %9715, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9717 = stablehlo.broadcast_in_dim %9716, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9718 = stablehlo.multiply %9717, %4556 : tensor<1x2048xf32> loc(#loc5262)
    %9719 = stablehlo.broadcast_in_dim %9689, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9720 = stablehlo.broadcast_in_dim %9719, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9721 = stablehlo.multiply %9720, %arg269 : tensor<1x2048xf32> loc(#loc5262)
    %9722 = stablehlo.add %9718, %9721 : tensor<1x2048xf32> loc(#loc5263)
    %9723 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9724 = stablehlo.broadcast_in_dim %9723, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9725 = stablehlo.subtract %9724, %9712 : tensor<1xf32> loc(#loc5260)
    %9726 = stablehlo.multiply %4556, %4556 : tensor<1x2048xf32> loc(#loc5264)
    %9727 = stablehlo.broadcast_in_dim %9725, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9728 = stablehlo.broadcast_in_dim %9727, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9729 = stablehlo.multiply %9728, %9726 : tensor<1x2048xf32> loc(#loc5262)
    %9730 = stablehlo.broadcast_in_dim %9712, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9731 = stablehlo.broadcast_in_dim %9730, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9732 = stablehlo.multiply %9731, %arg445 : tensor<1x2048xf32> loc(#loc5262)
    %9733 = stablehlo.add %9729, %9732 : tensor<1x2048xf32> loc(#loc5263)
    %9734 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9735 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9736 = stablehlo.broadcast_in_dim %9735, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9737 = stablehlo.add %9734, %9736 : tensor<1xf32> loc(#loc5259)
    %9738 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9739 = stablehlo.broadcast_in_dim %9738, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9740 = stablehlo.subtract %9737, %9739 : tensor<1xf32> loc(#loc5260)
    %9741 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9742 = stablehlo.broadcast_in_dim %9741, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9743 = stablehlo.power %9742, %9740 : tensor<1xf32> loc(#loc5261)
    %9744 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9745 = stablehlo.broadcast_in_dim %9744, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9746 = stablehlo.subtract %9745, %9743 : tensor<1xf32> loc(#loc5260)
    %9747 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9748 = stablehlo.broadcast_in_dim %9747, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9749 = stablehlo.multiply %9748, %9746 : tensor<1xf32> loc(#loc5236)
    %9750 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9751 = stablehlo.broadcast_in_dim %9750, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9752 = stablehlo.power %9751, %9737 : tensor<1xf32> loc(#loc5261)
    %9753 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9754 = stablehlo.broadcast_in_dim %9753, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9755 = stablehlo.subtract %9754, %9752 : tensor<1xf32> loc(#loc5260)
    %9756 = stablehlo.divide %9749, %9755 : tensor<1xf32> loc(#loc5253)
    %9757 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9758 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9759 = stablehlo.broadcast_in_dim %9758, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9760 = stablehlo.add %9757, %9759 : tensor<1xf32> loc(#loc5259)
    %9761 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9762 = stablehlo.broadcast_in_dim %9761, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9763 = stablehlo.subtract %9760, %9762 : tensor<1xf32> loc(#loc5260)
    %9764 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9765 = stablehlo.broadcast_in_dim %9764, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9766 = stablehlo.power %9765, %9763 : tensor<1xf32> loc(#loc5261)
    %9767 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9768 = stablehlo.broadcast_in_dim %9767, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9769 = stablehlo.subtract %9768, %9766 : tensor<1xf32> loc(#loc5260)
    %9770 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9771 = stablehlo.broadcast_in_dim %9770, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9772 = stablehlo.multiply %9771, %9769 : tensor<1xf32> loc(#loc5236)
    %9773 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9774 = stablehlo.broadcast_in_dim %9773, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9775 = stablehlo.power %9774, %9760 : tensor<1xf32> loc(#loc5261)
    %9776 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9777 = stablehlo.broadcast_in_dim %9776, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9778 = stablehlo.subtract %9777, %9775 : tensor<1xf32> loc(#loc5260)
    %9779 = stablehlo.divide %9772, %9778 : tensor<1xf32> loc(#loc5253)
    %9780 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9781 = stablehlo.broadcast_in_dim %9780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9782 = stablehlo.subtract %9781, %9756 : tensor<1xf32> loc(#loc5260)
    %9783 = stablehlo.broadcast_in_dim %9782, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9784 = stablehlo.broadcast_in_dim %9783, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9785 = stablehlo.multiply %9784, %4560 : tensor<1x2048xf32> loc(#loc5262)
    %9786 = stablehlo.broadcast_in_dim %9756, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9787 = stablehlo.broadcast_in_dim %9786, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9788 = stablehlo.multiply %9787, %arg270 : tensor<1x2048xf32> loc(#loc5262)
    %9789 = stablehlo.add %9785, %9788 : tensor<1x2048xf32> loc(#loc5263)
    %9790 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9791 = stablehlo.broadcast_in_dim %9790, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9792 = stablehlo.subtract %9791, %9779 : tensor<1xf32> loc(#loc5260)
    %9793 = stablehlo.multiply %4560, %4560 : tensor<1x2048xf32> loc(#loc5264)
    %9794 = stablehlo.broadcast_in_dim %9792, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9795 = stablehlo.broadcast_in_dim %9794, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9796 = stablehlo.multiply %9795, %9793 : tensor<1x2048xf32> loc(#loc5262)
    %9797 = stablehlo.broadcast_in_dim %9779, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9798 = stablehlo.broadcast_in_dim %9797, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9799 = stablehlo.multiply %9798, %arg446 : tensor<1x2048xf32> loc(#loc5262)
    %9800 = stablehlo.add %9796, %9799 : tensor<1x2048xf32> loc(#loc5263)
    %9801 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9802 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9803 = stablehlo.broadcast_in_dim %9802, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9804 = stablehlo.add %9801, %9803 : tensor<1xf32> loc(#loc5259)
    %9805 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9806 = stablehlo.broadcast_in_dim %9805, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9807 = stablehlo.subtract %9804, %9806 : tensor<1xf32> loc(#loc5260)
    %9808 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9809 = stablehlo.broadcast_in_dim %9808, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9810 = stablehlo.power %9809, %9807 : tensor<1xf32> loc(#loc5261)
    %9811 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9812 = stablehlo.broadcast_in_dim %9811, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9813 = stablehlo.subtract %9812, %9810 : tensor<1xf32> loc(#loc5260)
    %9814 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9815 = stablehlo.broadcast_in_dim %9814, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9816 = stablehlo.multiply %9815, %9813 : tensor<1xf32> loc(#loc5236)
    %9817 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9818 = stablehlo.broadcast_in_dim %9817, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9819 = stablehlo.power %9818, %9804 : tensor<1xf32> loc(#loc5261)
    %9820 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9821 = stablehlo.broadcast_in_dim %9820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9822 = stablehlo.subtract %9821, %9819 : tensor<1xf32> loc(#loc5260)
    %9823 = stablehlo.divide %9816, %9822 : tensor<1xf32> loc(#loc5253)
    %9824 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9825 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9826 = stablehlo.broadcast_in_dim %9825, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9827 = stablehlo.add %9824, %9826 : tensor<1xf32> loc(#loc5259)
    %9828 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9829 = stablehlo.broadcast_in_dim %9828, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9830 = stablehlo.subtract %9827, %9829 : tensor<1xf32> loc(#loc5260)
    %9831 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9832 = stablehlo.broadcast_in_dim %9831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9833 = stablehlo.power %9832, %9830 : tensor<1xf32> loc(#loc5261)
    %9834 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9835 = stablehlo.broadcast_in_dim %9834, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9836 = stablehlo.subtract %9835, %9833 : tensor<1xf32> loc(#loc5260)
    %9837 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9838 = stablehlo.broadcast_in_dim %9837, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9839 = stablehlo.multiply %9838, %9836 : tensor<1xf32> loc(#loc5236)
    %9840 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9841 = stablehlo.broadcast_in_dim %9840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9842 = stablehlo.power %9841, %9827 : tensor<1xf32> loc(#loc5261)
    %9843 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9844 = stablehlo.broadcast_in_dim %9843, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9845 = stablehlo.subtract %9844, %9842 : tensor<1xf32> loc(#loc5260)
    %9846 = stablehlo.divide %9839, %9845 : tensor<1xf32> loc(#loc5253)
    %9847 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9848 = stablehlo.broadcast_in_dim %9847, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9849 = stablehlo.subtract %9848, %9823 : tensor<1xf32> loc(#loc5260)
    %9850 = stablehlo.broadcast_in_dim %9849, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9851 = stablehlo.broadcast_in_dim %9850, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9852 = stablehlo.multiply %9851, %4564 : tensor<1x2048xf32> loc(#loc5262)
    %9853 = stablehlo.broadcast_in_dim %9823, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9854 = stablehlo.broadcast_in_dim %9853, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9855 = stablehlo.multiply %9854, %arg271 : tensor<1x2048xf32> loc(#loc5262)
    %9856 = stablehlo.add %9852, %9855 : tensor<1x2048xf32> loc(#loc5263)
    %9857 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9858 = stablehlo.broadcast_in_dim %9857, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9859 = stablehlo.subtract %9858, %9846 : tensor<1xf32> loc(#loc5260)
    %9860 = stablehlo.multiply %4564, %4564 : tensor<1x2048xf32> loc(#loc5264)
    %9861 = stablehlo.broadcast_in_dim %9859, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9862 = stablehlo.broadcast_in_dim %9861, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9863 = stablehlo.multiply %9862, %9860 : tensor<1x2048xf32> loc(#loc5262)
    %9864 = stablehlo.broadcast_in_dim %9846, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9865 = stablehlo.broadcast_in_dim %9864, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9866 = stablehlo.multiply %9865, %arg447 : tensor<1x2048xf32> loc(#loc5262)
    %9867 = stablehlo.add %9863, %9866 : tensor<1x2048xf32> loc(#loc5263)
    %9868 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9869 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9870 = stablehlo.broadcast_in_dim %9869, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9871 = stablehlo.add %9868, %9870 : tensor<1xf32> loc(#loc5259)
    %9872 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9873 = stablehlo.broadcast_in_dim %9872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9874 = stablehlo.subtract %9871, %9873 : tensor<1xf32> loc(#loc5260)
    %9875 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9876 = stablehlo.broadcast_in_dim %9875, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9877 = stablehlo.power %9876, %9874 : tensor<1xf32> loc(#loc5261)
    %9878 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9879 = stablehlo.broadcast_in_dim %9878, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9880 = stablehlo.subtract %9879, %9877 : tensor<1xf32> loc(#loc5260)
    %9881 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9882 = stablehlo.broadcast_in_dim %9881, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9883 = stablehlo.multiply %9882, %9880 : tensor<1xf32> loc(#loc5236)
    %9884 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9885 = stablehlo.broadcast_in_dim %9884, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9886 = stablehlo.power %9885, %9871 : tensor<1xf32> loc(#loc5261)
    %9887 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9888 = stablehlo.broadcast_in_dim %9887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9889 = stablehlo.subtract %9888, %9886 : tensor<1xf32> loc(#loc5260)
    %9890 = stablehlo.divide %9883, %9889 : tensor<1xf32> loc(#loc5253)
    %9891 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9892 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9893 = stablehlo.broadcast_in_dim %9892, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9894 = stablehlo.add %9891, %9893 : tensor<1xf32> loc(#loc5259)
    %9895 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9896 = stablehlo.broadcast_in_dim %9895, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9897 = stablehlo.subtract %9894, %9896 : tensor<1xf32> loc(#loc5260)
    %9898 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9899 = stablehlo.broadcast_in_dim %9898, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9900 = stablehlo.power %9899, %9897 : tensor<1xf32> loc(#loc5261)
    %9901 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9902 = stablehlo.broadcast_in_dim %9901, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9903 = stablehlo.subtract %9902, %9900 : tensor<1xf32> loc(#loc5260)
    %9904 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9905 = stablehlo.broadcast_in_dim %9904, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9906 = stablehlo.multiply %9905, %9903 : tensor<1xf32> loc(#loc5236)
    %9907 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9908 = stablehlo.broadcast_in_dim %9907, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9909 = stablehlo.power %9908, %9894 : tensor<1xf32> loc(#loc5261)
    %9910 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9911 = stablehlo.broadcast_in_dim %9910, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9912 = stablehlo.subtract %9911, %9909 : tensor<1xf32> loc(#loc5260)
    %9913 = stablehlo.divide %9906, %9912 : tensor<1xf32> loc(#loc5253)
    %9914 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9915 = stablehlo.broadcast_in_dim %9914, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9916 = stablehlo.subtract %9915, %9890 : tensor<1xf32> loc(#loc5260)
    %9917 = stablehlo.broadcast_in_dim %9916, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9918 = stablehlo.broadcast_in_dim %9917, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9919 = stablehlo.multiply %9918, %4568 : tensor<1x2048xf32> loc(#loc5262)
    %9920 = stablehlo.broadcast_in_dim %9890, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9921 = stablehlo.broadcast_in_dim %9920, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9922 = stablehlo.multiply %9921, %arg272 : tensor<1x2048xf32> loc(#loc5262)
    %9923 = stablehlo.add %9919, %9922 : tensor<1x2048xf32> loc(#loc5263)
    %9924 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9925 = stablehlo.broadcast_in_dim %9924, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9926 = stablehlo.subtract %9925, %9913 : tensor<1xf32> loc(#loc5260)
    %9927 = stablehlo.multiply %4568, %4568 : tensor<1x2048xf32> loc(#loc5264)
    %9928 = stablehlo.broadcast_in_dim %9926, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9929 = stablehlo.broadcast_in_dim %9928, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9930 = stablehlo.multiply %9929, %9927 : tensor<1x2048xf32> loc(#loc5262)
    %9931 = stablehlo.broadcast_in_dim %9913, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %9932 = stablehlo.broadcast_in_dim %9931, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %9933 = stablehlo.multiply %9932, %arg448 : tensor<1x2048xf32> loc(#loc5262)
    %9934 = stablehlo.add %9930, %9933 : tensor<1x2048xf32> loc(#loc5263)
    %9935 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9936 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9937 = stablehlo.broadcast_in_dim %9936, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9938 = stablehlo.add %9935, %9937 : tensor<1xf32> loc(#loc5259)
    %9939 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9940 = stablehlo.broadcast_in_dim %9939, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9941 = stablehlo.subtract %9938, %9940 : tensor<1xf32> loc(#loc5260)
    %9942 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9943 = stablehlo.broadcast_in_dim %9942, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9944 = stablehlo.power %9943, %9941 : tensor<1xf32> loc(#loc5261)
    %9945 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9946 = stablehlo.broadcast_in_dim %9945, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9947 = stablehlo.subtract %9946, %9944 : tensor<1xf32> loc(#loc5260)
    %9948 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9949 = stablehlo.broadcast_in_dim %9948, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9950 = stablehlo.multiply %9949, %9947 : tensor<1xf32> loc(#loc5236)
    %9951 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %9952 = stablehlo.broadcast_in_dim %9951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9953 = stablehlo.power %9952, %9938 : tensor<1xf32> loc(#loc5261)
    %9954 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9955 = stablehlo.broadcast_in_dim %9954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9956 = stablehlo.subtract %9955, %9953 : tensor<1xf32> loc(#loc5260)
    %9957 = stablehlo.divide %9950, %9956 : tensor<1xf32> loc(#loc5253)
    %9958 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %9959 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9960 = stablehlo.broadcast_in_dim %9959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %9961 = stablehlo.add %9958, %9960 : tensor<1xf32> loc(#loc5259)
    %9962 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9963 = stablehlo.broadcast_in_dim %9962, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9964 = stablehlo.subtract %9961, %9963 : tensor<1xf32> loc(#loc5260)
    %9965 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9966 = stablehlo.broadcast_in_dim %9965, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9967 = stablehlo.power %9966, %9964 : tensor<1xf32> loc(#loc5261)
    %9968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9969 = stablehlo.broadcast_in_dim %9968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9970 = stablehlo.subtract %9969, %9967 : tensor<1xf32> loc(#loc5260)
    %9971 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9972 = stablehlo.broadcast_in_dim %9971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %9973 = stablehlo.multiply %9972, %9970 : tensor<1xf32> loc(#loc5236)
    %9974 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %9975 = stablehlo.broadcast_in_dim %9974, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %9976 = stablehlo.power %9975, %9961 : tensor<1xf32> loc(#loc5261)
    %9977 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9978 = stablehlo.broadcast_in_dim %9977, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9979 = stablehlo.subtract %9978, %9976 : tensor<1xf32> loc(#loc5260)
    %9980 = stablehlo.divide %9973, %9979 : tensor<1xf32> loc(#loc5253)
    %9981 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9982 = stablehlo.broadcast_in_dim %9981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9983 = stablehlo.subtract %9982, %9957 : tensor<1xf32> loc(#loc5260)
    %9984 = stablehlo.broadcast_in_dim %9983, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9985 = stablehlo.broadcast_in_dim %9984, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9986 = stablehlo.multiply %9985, %4572 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9987 = stablehlo.broadcast_in_dim %9957, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9988 = stablehlo.broadcast_in_dim %9987, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9989 = stablehlo.multiply %9988, %arg273 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9990 = stablehlo.add %9986, %9989 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %9991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %9992 = stablehlo.broadcast_in_dim %9991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %9993 = stablehlo.subtract %9992, %9980 : tensor<1xf32> loc(#loc5260)
    %9994 = stablehlo.multiply %4572, %4572 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %9995 = stablehlo.broadcast_in_dim %9993, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9996 = stablehlo.broadcast_in_dim %9995, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9997 = stablehlo.multiply %9996, %9994 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %9998 = stablehlo.broadcast_in_dim %9980, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %9999 = stablehlo.broadcast_in_dim %9998, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10000 = stablehlo.multiply %9999, %arg449 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10001 = stablehlo.add %9997, %10000 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %10002 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10003 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10004 = stablehlo.broadcast_in_dim %10003, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10005 = stablehlo.add %10002, %10004 : tensor<1xf32> loc(#loc5259)
    %10006 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10007 = stablehlo.broadcast_in_dim %10006, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10008 = stablehlo.subtract %10005, %10007 : tensor<1xf32> loc(#loc5260)
    %10009 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10010 = stablehlo.broadcast_in_dim %10009, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10011 = stablehlo.power %10010, %10008 : tensor<1xf32> loc(#loc5261)
    %10012 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10013 = stablehlo.broadcast_in_dim %10012, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10014 = stablehlo.subtract %10013, %10011 : tensor<1xf32> loc(#loc5260)
    %10015 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10016 = stablehlo.broadcast_in_dim %10015, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10017 = stablehlo.multiply %10016, %10014 : tensor<1xf32> loc(#loc5236)
    %10018 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10019 = stablehlo.broadcast_in_dim %10018, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10020 = stablehlo.power %10019, %10005 : tensor<1xf32> loc(#loc5261)
    %10021 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10022 = stablehlo.broadcast_in_dim %10021, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10023 = stablehlo.subtract %10022, %10020 : tensor<1xf32> loc(#loc5260)
    %10024 = stablehlo.divide %10017, %10023 : tensor<1xf32> loc(#loc5253)
    %10025 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10026 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10027 = stablehlo.broadcast_in_dim %10026, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10028 = stablehlo.add %10025, %10027 : tensor<1xf32> loc(#loc5259)
    %10029 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10030 = stablehlo.broadcast_in_dim %10029, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10031 = stablehlo.subtract %10028, %10030 : tensor<1xf32> loc(#loc5260)
    %10032 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10033 = stablehlo.broadcast_in_dim %10032, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10034 = stablehlo.power %10033, %10031 : tensor<1xf32> loc(#loc5261)
    %10035 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10036 = stablehlo.broadcast_in_dim %10035, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10037 = stablehlo.subtract %10036, %10034 : tensor<1xf32> loc(#loc5260)
    %10038 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10039 = stablehlo.broadcast_in_dim %10038, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10040 = stablehlo.multiply %10039, %10037 : tensor<1xf32> loc(#loc5236)
    %10041 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10042 = stablehlo.broadcast_in_dim %10041, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10043 = stablehlo.power %10042, %10028 : tensor<1xf32> loc(#loc5261)
    %10044 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10045 = stablehlo.broadcast_in_dim %10044, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10046 = stablehlo.subtract %10045, %10043 : tensor<1xf32> loc(#loc5260)
    %10047 = stablehlo.divide %10040, %10046 : tensor<1xf32> loc(#loc5253)
    %10048 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10049 = stablehlo.broadcast_in_dim %10048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10050 = stablehlo.subtract %10049, %10024 : tensor<1xf32> loc(#loc5260)
    %10051 = stablehlo.broadcast_in_dim %10050, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10052 = stablehlo.broadcast_in_dim %10051, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10053 = stablehlo.multiply %10052, %4576 : tensor<1x64xf32> loc(#loc5262)
    %10054 = stablehlo.broadcast_in_dim %10024, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10055 = stablehlo.broadcast_in_dim %10054, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10056 = stablehlo.multiply %10055, %arg274 : tensor<1x64xf32> loc(#loc5262)
    %10057 = stablehlo.add %10053, %10056 : tensor<1x64xf32> loc(#loc5263)
    %10058 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10059 = stablehlo.broadcast_in_dim %10058, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10060 = stablehlo.subtract %10059, %10047 : tensor<1xf32> loc(#loc5260)
    %10061 = stablehlo.multiply %4576, %4576 : tensor<1x64xf32> loc(#loc5264)
    %10062 = stablehlo.broadcast_in_dim %10060, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10063 = stablehlo.broadcast_in_dim %10062, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10064 = stablehlo.multiply %10063, %10061 : tensor<1x64xf32> loc(#loc5262)
    %10065 = stablehlo.broadcast_in_dim %10047, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10066 = stablehlo.broadcast_in_dim %10065, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10067 = stablehlo.multiply %10066, %arg450 : tensor<1x64xf32> loc(#loc5262)
    %10068 = stablehlo.add %10064, %10067 : tensor<1x64xf32> loc(#loc5263)
    %10069 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10070 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10071 = stablehlo.broadcast_in_dim %10070, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10072 = stablehlo.add %10069, %10071 : tensor<1xf32> loc(#loc5259)
    %10073 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10074 = stablehlo.broadcast_in_dim %10073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10075 = stablehlo.subtract %10072, %10074 : tensor<1xf32> loc(#loc5260)
    %10076 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10077 = stablehlo.broadcast_in_dim %10076, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10078 = stablehlo.power %10077, %10075 : tensor<1xf32> loc(#loc5261)
    %10079 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10080 = stablehlo.broadcast_in_dim %10079, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10081 = stablehlo.subtract %10080, %10078 : tensor<1xf32> loc(#loc5260)
    %10082 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10083 = stablehlo.broadcast_in_dim %10082, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10084 = stablehlo.multiply %10083, %10081 : tensor<1xf32> loc(#loc5236)
    %10085 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10086 = stablehlo.broadcast_in_dim %10085, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10087 = stablehlo.power %10086, %10072 : tensor<1xf32> loc(#loc5261)
    %10088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10089 = stablehlo.broadcast_in_dim %10088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10090 = stablehlo.subtract %10089, %10087 : tensor<1xf32> loc(#loc5260)
    %10091 = stablehlo.divide %10084, %10090 : tensor<1xf32> loc(#loc5253)
    %10092 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10093 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10094 = stablehlo.broadcast_in_dim %10093, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10095 = stablehlo.add %10092, %10094 : tensor<1xf32> loc(#loc5259)
    %10096 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10097 = stablehlo.broadcast_in_dim %10096, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10098 = stablehlo.subtract %10095, %10097 : tensor<1xf32> loc(#loc5260)
    %10099 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10100 = stablehlo.broadcast_in_dim %10099, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10101 = stablehlo.power %10100, %10098 : tensor<1xf32> loc(#loc5261)
    %10102 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10103 = stablehlo.broadcast_in_dim %10102, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10104 = stablehlo.subtract %10103, %10101 : tensor<1xf32> loc(#loc5260)
    %10105 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10106 = stablehlo.broadcast_in_dim %10105, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10107 = stablehlo.multiply %10106, %10104 : tensor<1xf32> loc(#loc5236)
    %10108 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10109 = stablehlo.broadcast_in_dim %10108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10110 = stablehlo.power %10109, %10095 : tensor<1xf32> loc(#loc5261)
    %10111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10112 = stablehlo.broadcast_in_dim %10111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10113 = stablehlo.subtract %10112, %10110 : tensor<1xf32> loc(#loc5260)
    %10114 = stablehlo.divide %10107, %10113 : tensor<1xf32> loc(#loc5253)
    %10115 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10116 = stablehlo.broadcast_in_dim %10115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10117 = stablehlo.subtract %10116, %10091 : tensor<1xf32> loc(#loc5260)
    %10118 = stablehlo.broadcast_in_dim %10117, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10119 = stablehlo.broadcast_in_dim %10118, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10120 = stablehlo.multiply %10119, %4580 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10121 = stablehlo.broadcast_in_dim %10091, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10122 = stablehlo.broadcast_in_dim %10121, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10123 = stablehlo.multiply %10122, %arg275 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10124 = stablehlo.add %10120, %10123 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %10125 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10126 = stablehlo.broadcast_in_dim %10125, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10127 = stablehlo.subtract %10126, %10114 : tensor<1xf32> loc(#loc5260)
    %10128 = stablehlo.multiply %4580, %4580 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %10129 = stablehlo.broadcast_in_dim %10127, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10130 = stablehlo.broadcast_in_dim %10129, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10131 = stablehlo.multiply %10130, %10128 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10132 = stablehlo.broadcast_in_dim %10114, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10133 = stablehlo.broadcast_in_dim %10132, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10134 = stablehlo.multiply %10133, %arg451 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10135 = stablehlo.add %10131, %10134 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %10136 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10137 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10138 = stablehlo.broadcast_in_dim %10137, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10139 = stablehlo.add %10136, %10138 : tensor<1xf32> loc(#loc5259)
    %10140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10141 = stablehlo.broadcast_in_dim %10140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10142 = stablehlo.subtract %10139, %10141 : tensor<1xf32> loc(#loc5260)
    %10143 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10144 = stablehlo.broadcast_in_dim %10143, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10145 = stablehlo.power %10144, %10142 : tensor<1xf32> loc(#loc5261)
    %10146 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10147 = stablehlo.broadcast_in_dim %10146, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10148 = stablehlo.subtract %10147, %10145 : tensor<1xf32> loc(#loc5260)
    %10149 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10150 = stablehlo.broadcast_in_dim %10149, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10151 = stablehlo.multiply %10150, %10148 : tensor<1xf32> loc(#loc5236)
    %10152 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10153 = stablehlo.broadcast_in_dim %10152, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10154 = stablehlo.power %10153, %10139 : tensor<1xf32> loc(#loc5261)
    %10155 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10156 = stablehlo.broadcast_in_dim %10155, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10157 = stablehlo.subtract %10156, %10154 : tensor<1xf32> loc(#loc5260)
    %10158 = stablehlo.divide %10151, %10157 : tensor<1xf32> loc(#loc5253)
    %10159 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10160 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10161 = stablehlo.broadcast_in_dim %10160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10162 = stablehlo.add %10159, %10161 : tensor<1xf32> loc(#loc5259)
    %10163 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10164 = stablehlo.broadcast_in_dim %10163, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10165 = stablehlo.subtract %10162, %10164 : tensor<1xf32> loc(#loc5260)
    %10166 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10167 = stablehlo.broadcast_in_dim %10166, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10168 = stablehlo.power %10167, %10165 : tensor<1xf32> loc(#loc5261)
    %10169 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10170 = stablehlo.broadcast_in_dim %10169, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10171 = stablehlo.subtract %10170, %10168 : tensor<1xf32> loc(#loc5260)
    %10172 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10173 = stablehlo.broadcast_in_dim %10172, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10174 = stablehlo.multiply %10173, %10171 : tensor<1xf32> loc(#loc5236)
    %10175 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10176 = stablehlo.broadcast_in_dim %10175, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10177 = stablehlo.power %10176, %10162 : tensor<1xf32> loc(#loc5261)
    %10178 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10179 = stablehlo.broadcast_in_dim %10178, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10180 = stablehlo.subtract %10179, %10177 : tensor<1xf32> loc(#loc5260)
    %10181 = stablehlo.divide %10174, %10180 : tensor<1xf32> loc(#loc5253)
    %10182 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10183 = stablehlo.broadcast_in_dim %10182, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10184 = stablehlo.subtract %10183, %10158 : tensor<1xf32> loc(#loc5260)
    %10185 = stablehlo.broadcast_in_dim %10184, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10186 = stablehlo.broadcast_in_dim %10185, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10187 = stablehlo.multiply %10186, %4584 : tensor<1x8192xf32> loc(#loc5262)
    %10188 = stablehlo.broadcast_in_dim %10158, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10189 = stablehlo.broadcast_in_dim %10188, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10190 = stablehlo.multiply %10189, %arg276 : tensor<1x8192xf32> loc(#loc5262)
    %10191 = stablehlo.add %10187, %10190 : tensor<1x8192xf32> loc(#loc5263)
    %10192 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10193 = stablehlo.broadcast_in_dim %10192, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10194 = stablehlo.subtract %10193, %10181 : tensor<1xf32> loc(#loc5260)
    %10195 = stablehlo.multiply %4584, %4584 : tensor<1x8192xf32> loc(#loc5264)
    %10196 = stablehlo.broadcast_in_dim %10194, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10197 = stablehlo.broadcast_in_dim %10196, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10198 = stablehlo.multiply %10197, %10195 : tensor<1x8192xf32> loc(#loc5262)
    %10199 = stablehlo.broadcast_in_dim %10181, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10200 = stablehlo.broadcast_in_dim %10199, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10201 = stablehlo.multiply %10200, %arg452 : tensor<1x8192xf32> loc(#loc5262)
    %10202 = stablehlo.add %10198, %10201 : tensor<1x8192xf32> loc(#loc5263)
    %10203 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10204 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10205 = stablehlo.broadcast_in_dim %10204, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10206 = stablehlo.add %10203, %10205 : tensor<1xf32> loc(#loc5259)
    %10207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10208 = stablehlo.broadcast_in_dim %10207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10209 = stablehlo.subtract %10206, %10208 : tensor<1xf32> loc(#loc5260)
    %10210 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10211 = stablehlo.broadcast_in_dim %10210, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10212 = stablehlo.power %10211, %10209 : tensor<1xf32> loc(#loc5261)
    %10213 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10214 = stablehlo.broadcast_in_dim %10213, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10215 = stablehlo.subtract %10214, %10212 : tensor<1xf32> loc(#loc5260)
    %10216 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10217 = stablehlo.broadcast_in_dim %10216, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10218 = stablehlo.multiply %10217, %10215 : tensor<1xf32> loc(#loc5236)
    %10219 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10220 = stablehlo.broadcast_in_dim %10219, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10221 = stablehlo.power %10220, %10206 : tensor<1xf32> loc(#loc5261)
    %10222 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10223 = stablehlo.broadcast_in_dim %10222, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10224 = stablehlo.subtract %10223, %10221 : tensor<1xf32> loc(#loc5260)
    %10225 = stablehlo.divide %10218, %10224 : tensor<1xf32> loc(#loc5253)
    %10226 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10227 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10228 = stablehlo.broadcast_in_dim %10227, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10229 = stablehlo.add %10226, %10228 : tensor<1xf32> loc(#loc5259)
    %10230 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10231 = stablehlo.broadcast_in_dim %10230, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10232 = stablehlo.subtract %10229, %10231 : tensor<1xf32> loc(#loc5260)
    %10233 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10234 = stablehlo.broadcast_in_dim %10233, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10235 = stablehlo.power %10234, %10232 : tensor<1xf32> loc(#loc5261)
    %10236 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10237 = stablehlo.broadcast_in_dim %10236, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10238 = stablehlo.subtract %10237, %10235 : tensor<1xf32> loc(#loc5260)
    %10239 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10240 = stablehlo.broadcast_in_dim %10239, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10241 = stablehlo.multiply %10240, %10238 : tensor<1xf32> loc(#loc5236)
    %10242 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10243 = stablehlo.broadcast_in_dim %10242, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10244 = stablehlo.power %10243, %10229 : tensor<1xf32> loc(#loc5261)
    %10245 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10246 = stablehlo.broadcast_in_dim %10245, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10247 = stablehlo.subtract %10246, %10244 : tensor<1xf32> loc(#loc5260)
    %10248 = stablehlo.divide %10241, %10247 : tensor<1xf32> loc(#loc5253)
    %10249 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10250 = stablehlo.broadcast_in_dim %10249, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10251 = stablehlo.subtract %10250, %10225 : tensor<1xf32> loc(#loc5260)
    %10252 = stablehlo.broadcast_in_dim %10251, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10253 = stablehlo.broadcast_in_dim %10252, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %10254 = stablehlo.multiply %10253, %4588 : tensor<1x2048x8192xf32> loc(#loc5262)
    %10255 = stablehlo.broadcast_in_dim %10225, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10256 = stablehlo.broadcast_in_dim %10255, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %10257 = stablehlo.multiply %10256, %arg277 : tensor<1x2048x8192xf32> loc(#loc5262)
    %10258 = stablehlo.add %10254, %10257 : tensor<1x2048x8192xf32> loc(#loc5263)
    %10259 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10260 = stablehlo.broadcast_in_dim %10259, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10261 = stablehlo.subtract %10260, %10248 : tensor<1xf32> loc(#loc5260)
    %10262 = stablehlo.multiply %4588, %4588 : tensor<1x2048x8192xf32> loc(#loc5264)
    %10263 = stablehlo.broadcast_in_dim %10261, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10264 = stablehlo.broadcast_in_dim %10263, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %10265 = stablehlo.multiply %10264, %10262 : tensor<1x2048x8192xf32> loc(#loc5262)
    %10266 = stablehlo.broadcast_in_dim %10248, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10267 = stablehlo.broadcast_in_dim %10266, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %10268 = stablehlo.multiply %10267, %arg453 : tensor<1x2048x8192xf32> loc(#loc5262)
    %10269 = stablehlo.add %10265, %10268 : tensor<1x2048x8192xf32> loc(#loc5263)
    %10270 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10272 = stablehlo.broadcast_in_dim %10271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10273 = stablehlo.add %10270, %10272 : tensor<1xf32> loc(#loc5259)
    %10274 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10275 = stablehlo.broadcast_in_dim %10274, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10276 = stablehlo.subtract %10273, %10275 : tensor<1xf32> loc(#loc5260)
    %10277 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10278 = stablehlo.broadcast_in_dim %10277, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10279 = stablehlo.power %10278, %10276 : tensor<1xf32> loc(#loc5261)
    %10280 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10281 = stablehlo.broadcast_in_dim %10280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10282 = stablehlo.subtract %10281, %10279 : tensor<1xf32> loc(#loc5260)
    %10283 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10284 = stablehlo.broadcast_in_dim %10283, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10285 = stablehlo.multiply %10284, %10282 : tensor<1xf32> loc(#loc5236)
    %10286 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10287 = stablehlo.broadcast_in_dim %10286, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10288 = stablehlo.power %10287, %10273 : tensor<1xf32> loc(#loc5261)
    %10289 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10290 = stablehlo.broadcast_in_dim %10289, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10291 = stablehlo.subtract %10290, %10288 : tensor<1xf32> loc(#loc5260)
    %10292 = stablehlo.divide %10285, %10291 : tensor<1xf32> loc(#loc5253)
    %10293 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10294 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10295 = stablehlo.broadcast_in_dim %10294, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10296 = stablehlo.add %10293, %10295 : tensor<1xf32> loc(#loc5259)
    %10297 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10298 = stablehlo.broadcast_in_dim %10297, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10299 = stablehlo.subtract %10296, %10298 : tensor<1xf32> loc(#loc5260)
    %10300 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10301 = stablehlo.broadcast_in_dim %10300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10302 = stablehlo.power %10301, %10299 : tensor<1xf32> loc(#loc5261)
    %10303 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10304 = stablehlo.broadcast_in_dim %10303, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10305 = stablehlo.subtract %10304, %10302 : tensor<1xf32> loc(#loc5260)
    %10306 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10307 = stablehlo.broadcast_in_dim %10306, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10308 = stablehlo.multiply %10307, %10305 : tensor<1xf32> loc(#loc5236)
    %10309 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10310 = stablehlo.broadcast_in_dim %10309, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10311 = stablehlo.power %10310, %10296 : tensor<1xf32> loc(#loc5261)
    %10312 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10313 = stablehlo.broadcast_in_dim %10312, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10314 = stablehlo.subtract %10313, %10311 : tensor<1xf32> loc(#loc5260)
    %10315 = stablehlo.divide %10308, %10314 : tensor<1xf32> loc(#loc5253)
    %10316 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10317 = stablehlo.broadcast_in_dim %10316, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10318 = stablehlo.subtract %10317, %10292 : tensor<1xf32> loc(#loc5260)
    %10319 = stablehlo.broadcast_in_dim %10318, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10320 = stablehlo.broadcast_in_dim %10319, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10321 = stablehlo.multiply %10320, %4592 : tensor<1x2048xf32> loc(#loc5262)
    %10322 = stablehlo.broadcast_in_dim %10292, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10323 = stablehlo.broadcast_in_dim %10322, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10324 = stablehlo.multiply %10323, %arg278 : tensor<1x2048xf32> loc(#loc5262)
    %10325 = stablehlo.add %10321, %10324 : tensor<1x2048xf32> loc(#loc5263)
    %10326 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10327 = stablehlo.broadcast_in_dim %10326, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10328 = stablehlo.subtract %10327, %10315 : tensor<1xf32> loc(#loc5260)
    %10329 = stablehlo.multiply %4592, %4592 : tensor<1x2048xf32> loc(#loc5264)
    %10330 = stablehlo.broadcast_in_dim %10328, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10331 = stablehlo.broadcast_in_dim %10330, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10332 = stablehlo.multiply %10331, %10329 : tensor<1x2048xf32> loc(#loc5262)
    %10333 = stablehlo.broadcast_in_dim %10315, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10334 = stablehlo.broadcast_in_dim %10333, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10335 = stablehlo.multiply %10334, %arg454 : tensor<1x2048xf32> loc(#loc5262)
    %10336 = stablehlo.add %10332, %10335 : tensor<1x2048xf32> loc(#loc5263)
    %10337 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10338 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10339 = stablehlo.broadcast_in_dim %10338, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10340 = stablehlo.add %10337, %10339 : tensor<1xf32> loc(#loc5259)
    %10341 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10342 = stablehlo.broadcast_in_dim %10341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10343 = stablehlo.subtract %10340, %10342 : tensor<1xf32> loc(#loc5260)
    %10344 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10345 = stablehlo.broadcast_in_dim %10344, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10346 = stablehlo.power %10345, %10343 : tensor<1xf32> loc(#loc5261)
    %10347 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10348 = stablehlo.broadcast_in_dim %10347, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10349 = stablehlo.subtract %10348, %10346 : tensor<1xf32> loc(#loc5260)
    %10350 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10351 = stablehlo.broadcast_in_dim %10350, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10352 = stablehlo.multiply %10351, %10349 : tensor<1xf32> loc(#loc5236)
    %10353 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10354 = stablehlo.broadcast_in_dim %10353, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10355 = stablehlo.power %10354, %10340 : tensor<1xf32> loc(#loc5261)
    %10356 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10357 = stablehlo.broadcast_in_dim %10356, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10358 = stablehlo.subtract %10357, %10355 : tensor<1xf32> loc(#loc5260)
    %10359 = stablehlo.divide %10352, %10358 : tensor<1xf32> loc(#loc5253)
    %10360 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10361 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10362 = stablehlo.broadcast_in_dim %10361, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10363 = stablehlo.add %10360, %10362 : tensor<1xf32> loc(#loc5259)
    %10364 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10365 = stablehlo.broadcast_in_dim %10364, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10366 = stablehlo.subtract %10363, %10365 : tensor<1xf32> loc(#loc5260)
    %10367 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10368 = stablehlo.broadcast_in_dim %10367, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10369 = stablehlo.power %10368, %10366 : tensor<1xf32> loc(#loc5261)
    %10370 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10371 = stablehlo.broadcast_in_dim %10370, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10372 = stablehlo.subtract %10371, %10369 : tensor<1xf32> loc(#loc5260)
    %10373 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10374 = stablehlo.broadcast_in_dim %10373, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10375 = stablehlo.multiply %10374, %10372 : tensor<1xf32> loc(#loc5236)
    %10376 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10377 = stablehlo.broadcast_in_dim %10376, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10378 = stablehlo.power %10377, %10363 : tensor<1xf32> loc(#loc5261)
    %10379 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10380 = stablehlo.broadcast_in_dim %10379, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10381 = stablehlo.subtract %10380, %10378 : tensor<1xf32> loc(#loc5260)
    %10382 = stablehlo.divide %10375, %10381 : tensor<1xf32> loc(#loc5253)
    %10383 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10384 = stablehlo.broadcast_in_dim %10383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10385 = stablehlo.subtract %10384, %10359 : tensor<1xf32> loc(#loc5260)
    %10386 = stablehlo.broadcast_in_dim %10385, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10387 = stablehlo.broadcast_in_dim %10386, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %10388 = stablehlo.multiply %10387, %4596 : tensor<1x8192x2048xf32> loc(#loc5262)
    %10389 = stablehlo.broadcast_in_dim %10359, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10390 = stablehlo.broadcast_in_dim %10389, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %10391 = stablehlo.multiply %10390, %arg279 : tensor<1x8192x2048xf32> loc(#loc5262)
    %10392 = stablehlo.add %10388, %10391 : tensor<1x8192x2048xf32> loc(#loc5263)
    %10393 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10394 = stablehlo.broadcast_in_dim %10393, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10395 = stablehlo.subtract %10394, %10382 : tensor<1xf32> loc(#loc5260)
    %10396 = stablehlo.multiply %4596, %4596 : tensor<1x8192x2048xf32> loc(#loc5264)
    %10397 = stablehlo.broadcast_in_dim %10395, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10398 = stablehlo.broadcast_in_dim %10397, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %10399 = stablehlo.multiply %10398, %10396 : tensor<1x8192x2048xf32> loc(#loc5262)
    %10400 = stablehlo.broadcast_in_dim %10382, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10401 = stablehlo.broadcast_in_dim %10400, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %10402 = stablehlo.multiply %10401, %arg455 : tensor<1x8192x2048xf32> loc(#loc5262)
    %10403 = stablehlo.add %10399, %10402 : tensor<1x8192x2048xf32> loc(#loc5263)
    %10404 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10405 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10406 = stablehlo.broadcast_in_dim %10405, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10407 = stablehlo.add %10404, %10406 : tensor<1xf32> loc(#loc5259)
    %10408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10409 = stablehlo.broadcast_in_dim %10408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10410 = stablehlo.subtract %10407, %10409 : tensor<1xf32> loc(#loc5260)
    %10411 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10412 = stablehlo.broadcast_in_dim %10411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10413 = stablehlo.power %10412, %10410 : tensor<1xf32> loc(#loc5261)
    %10414 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10415 = stablehlo.broadcast_in_dim %10414, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10416 = stablehlo.subtract %10415, %10413 : tensor<1xf32> loc(#loc5260)
    %10417 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10418 = stablehlo.broadcast_in_dim %10417, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10419 = stablehlo.multiply %10418, %10416 : tensor<1xf32> loc(#loc5236)
    %10420 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10421 = stablehlo.broadcast_in_dim %10420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10422 = stablehlo.power %10421, %10407 : tensor<1xf32> loc(#loc5261)
    %10423 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10424 = stablehlo.broadcast_in_dim %10423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10425 = stablehlo.subtract %10424, %10422 : tensor<1xf32> loc(#loc5260)
    %10426 = stablehlo.divide %10419, %10425 : tensor<1xf32> loc(#loc5253)
    %10427 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10429 = stablehlo.broadcast_in_dim %10428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10430 = stablehlo.add %10427, %10429 : tensor<1xf32> loc(#loc5259)
    %10431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10432 = stablehlo.broadcast_in_dim %10431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10433 = stablehlo.subtract %10430, %10432 : tensor<1xf32> loc(#loc5260)
    %10434 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10435 = stablehlo.broadcast_in_dim %10434, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10436 = stablehlo.power %10435, %10433 : tensor<1xf32> loc(#loc5261)
    %10437 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10438 = stablehlo.broadcast_in_dim %10437, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10439 = stablehlo.subtract %10438, %10436 : tensor<1xf32> loc(#loc5260)
    %10440 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10441 = stablehlo.broadcast_in_dim %10440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10442 = stablehlo.multiply %10441, %10439 : tensor<1xf32> loc(#loc5236)
    %10443 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10444 = stablehlo.broadcast_in_dim %10443, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10445 = stablehlo.power %10444, %10430 : tensor<1xf32> loc(#loc5261)
    %10446 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10447 = stablehlo.broadcast_in_dim %10446, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10448 = stablehlo.subtract %10447, %10445 : tensor<1xf32> loc(#loc5260)
    %10449 = stablehlo.divide %10442, %10448 : tensor<1xf32> loc(#loc5253)
    %10450 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10451 = stablehlo.broadcast_in_dim %10450, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10452 = stablehlo.subtract %10451, %10426 : tensor<1xf32> loc(#loc5260)
    %10453 = stablehlo.broadcast_in_dim %10452, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10454 = stablehlo.broadcast_in_dim %10453, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10455 = stablehlo.multiply %10454, %4600 : tensor<1x2048xf32> loc(#loc5262)
    %10456 = stablehlo.broadcast_in_dim %10426, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10457 = stablehlo.broadcast_in_dim %10456, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10458 = stablehlo.multiply %10457, %arg280 : tensor<1x2048xf32> loc(#loc5262)
    %10459 = stablehlo.add %10455, %10458 : tensor<1x2048xf32> loc(#loc5263)
    %10460 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10461 = stablehlo.broadcast_in_dim %10460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10462 = stablehlo.subtract %10461, %10449 : tensor<1xf32> loc(#loc5260)
    %10463 = stablehlo.multiply %4600, %4600 : tensor<1x2048xf32> loc(#loc5264)
    %10464 = stablehlo.broadcast_in_dim %10462, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10465 = stablehlo.broadcast_in_dim %10464, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10466 = stablehlo.multiply %10465, %10463 : tensor<1x2048xf32> loc(#loc5262)
    %10467 = stablehlo.broadcast_in_dim %10449, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10468 = stablehlo.broadcast_in_dim %10467, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10469 = stablehlo.multiply %10468, %arg456 : tensor<1x2048xf32> loc(#loc5262)
    %10470 = stablehlo.add %10466, %10469 : tensor<1x2048xf32> loc(#loc5263)
    %10471 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10472 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10473 = stablehlo.broadcast_in_dim %10472, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10474 = stablehlo.add %10471, %10473 : tensor<1xf32> loc(#loc5259)
    %10475 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10476 = stablehlo.broadcast_in_dim %10475, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10477 = stablehlo.subtract %10474, %10476 : tensor<1xf32> loc(#loc5260)
    %10478 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10479 = stablehlo.broadcast_in_dim %10478, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10480 = stablehlo.power %10479, %10477 : tensor<1xf32> loc(#loc5261)
    %10481 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10482 = stablehlo.broadcast_in_dim %10481, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10483 = stablehlo.subtract %10482, %10480 : tensor<1xf32> loc(#loc5260)
    %10484 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10485 = stablehlo.broadcast_in_dim %10484, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10486 = stablehlo.multiply %10485, %10483 : tensor<1xf32> loc(#loc5236)
    %10487 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10488 = stablehlo.broadcast_in_dim %10487, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10489 = stablehlo.power %10488, %10474 : tensor<1xf32> loc(#loc5261)
    %10490 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10491 = stablehlo.broadcast_in_dim %10490, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10492 = stablehlo.subtract %10491, %10489 : tensor<1xf32> loc(#loc5260)
    %10493 = stablehlo.divide %10486, %10492 : tensor<1xf32> loc(#loc5253)
    %10494 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10495 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10496 = stablehlo.broadcast_in_dim %10495, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10497 = stablehlo.add %10494, %10496 : tensor<1xf32> loc(#loc5259)
    %10498 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10499 = stablehlo.broadcast_in_dim %10498, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10500 = stablehlo.subtract %10497, %10499 : tensor<1xf32> loc(#loc5260)
    %10501 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10502 = stablehlo.broadcast_in_dim %10501, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10503 = stablehlo.power %10502, %10500 : tensor<1xf32> loc(#loc5261)
    %10504 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10505 = stablehlo.broadcast_in_dim %10504, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10506 = stablehlo.subtract %10505, %10503 : tensor<1xf32> loc(#loc5260)
    %10507 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10508 = stablehlo.broadcast_in_dim %10507, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10509 = stablehlo.multiply %10508, %10506 : tensor<1xf32> loc(#loc5236)
    %10510 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10511 = stablehlo.broadcast_in_dim %10510, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10512 = stablehlo.power %10511, %10497 : tensor<1xf32> loc(#loc5261)
    %10513 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10514 = stablehlo.broadcast_in_dim %10513, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10515 = stablehlo.subtract %10514, %10512 : tensor<1xf32> loc(#loc5260)
    %10516 = stablehlo.divide %10509, %10515 : tensor<1xf32> loc(#loc5253)
    %10517 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10518 = stablehlo.broadcast_in_dim %10517, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10519 = stablehlo.subtract %10518, %10493 : tensor<1xf32> loc(#loc5260)
    %10520 = stablehlo.broadcast_in_dim %10519, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10521 = stablehlo.broadcast_in_dim %10520, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10522 = stablehlo.multiply %10521, %4604 : tensor<1x2048xf32> loc(#loc5262)
    %10523 = stablehlo.broadcast_in_dim %10493, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10524 = stablehlo.broadcast_in_dim %10523, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10525 = stablehlo.multiply %10524, %arg281 : tensor<1x2048xf32> loc(#loc5262)
    %10526 = stablehlo.add %10522, %10525 : tensor<1x2048xf32> loc(#loc5263)
    %10527 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10528 = stablehlo.broadcast_in_dim %10527, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10529 = stablehlo.subtract %10528, %10516 : tensor<1xf32> loc(#loc5260)
    %10530 = stablehlo.multiply %4604, %4604 : tensor<1x2048xf32> loc(#loc5264)
    %10531 = stablehlo.broadcast_in_dim %10529, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10532 = stablehlo.broadcast_in_dim %10531, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10533 = stablehlo.multiply %10532, %10530 : tensor<1x2048xf32> loc(#loc5262)
    %10534 = stablehlo.broadcast_in_dim %10516, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10535 = stablehlo.broadcast_in_dim %10534, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10536 = stablehlo.multiply %10535, %arg457 : tensor<1x2048xf32> loc(#loc5262)
    %10537 = stablehlo.add %10533, %10536 : tensor<1x2048xf32> loc(#loc5263)
    %10538 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10539 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10540 = stablehlo.broadcast_in_dim %10539, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10541 = stablehlo.add %10538, %10540 : tensor<1xf32> loc(#loc5259)
    %10542 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10543 = stablehlo.broadcast_in_dim %10542, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10544 = stablehlo.subtract %10541, %10543 : tensor<1xf32> loc(#loc5260)
    %10545 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10546 = stablehlo.broadcast_in_dim %10545, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10547 = stablehlo.power %10546, %10544 : tensor<1xf32> loc(#loc5261)
    %10548 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10549 = stablehlo.broadcast_in_dim %10548, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10550 = stablehlo.subtract %10549, %10547 : tensor<1xf32> loc(#loc5260)
    %10551 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10552 = stablehlo.broadcast_in_dim %10551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10553 = stablehlo.multiply %10552, %10550 : tensor<1xf32> loc(#loc5236)
    %10554 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10555 = stablehlo.broadcast_in_dim %10554, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10556 = stablehlo.power %10555, %10541 : tensor<1xf32> loc(#loc5261)
    %10557 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10558 = stablehlo.broadcast_in_dim %10557, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10559 = stablehlo.subtract %10558, %10556 : tensor<1xf32> loc(#loc5260)
    %10560 = stablehlo.divide %10553, %10559 : tensor<1xf32> loc(#loc5253)
    %10561 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10562 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10563 = stablehlo.broadcast_in_dim %10562, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10564 = stablehlo.add %10561, %10563 : tensor<1xf32> loc(#loc5259)
    %10565 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10566 = stablehlo.broadcast_in_dim %10565, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10567 = stablehlo.subtract %10564, %10566 : tensor<1xf32> loc(#loc5260)
    %10568 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10569 = stablehlo.broadcast_in_dim %10568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10570 = stablehlo.power %10569, %10567 : tensor<1xf32> loc(#loc5261)
    %10571 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10572 = stablehlo.broadcast_in_dim %10571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10573 = stablehlo.subtract %10572, %10570 : tensor<1xf32> loc(#loc5260)
    %10574 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10575 = stablehlo.broadcast_in_dim %10574, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10576 = stablehlo.multiply %10575, %10573 : tensor<1xf32> loc(#loc5236)
    %10577 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10578 = stablehlo.broadcast_in_dim %10577, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10579 = stablehlo.power %10578, %10564 : tensor<1xf32> loc(#loc5261)
    %10580 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10581 = stablehlo.broadcast_in_dim %10580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10582 = stablehlo.subtract %10581, %10579 : tensor<1xf32> loc(#loc5260)
    %10583 = stablehlo.divide %10576, %10582 : tensor<1xf32> loc(#loc5253)
    %10584 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10585 = stablehlo.broadcast_in_dim %10584, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10586 = stablehlo.subtract %10585, %10560 : tensor<1xf32> loc(#loc5260)
    %10587 = stablehlo.broadcast_in_dim %10586, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10588 = stablehlo.broadcast_in_dim %10587, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10589 = stablehlo.multiply %10588, %4608 : tensor<1x2048xf32> loc(#loc5262)
    %10590 = stablehlo.broadcast_in_dim %10560, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10591 = stablehlo.broadcast_in_dim %10590, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10592 = stablehlo.multiply %10591, %arg282 : tensor<1x2048xf32> loc(#loc5262)
    %10593 = stablehlo.add %10589, %10592 : tensor<1x2048xf32> loc(#loc5263)
    %10594 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10595 = stablehlo.broadcast_in_dim %10594, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10596 = stablehlo.subtract %10595, %10583 : tensor<1xf32> loc(#loc5260)
    %10597 = stablehlo.multiply %4608, %4608 : tensor<1x2048xf32> loc(#loc5264)
    %10598 = stablehlo.broadcast_in_dim %10596, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10599 = stablehlo.broadcast_in_dim %10598, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10600 = stablehlo.multiply %10599, %10597 : tensor<1x2048xf32> loc(#loc5262)
    %10601 = stablehlo.broadcast_in_dim %10583, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10602 = stablehlo.broadcast_in_dim %10601, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10603 = stablehlo.multiply %10602, %arg458 : tensor<1x2048xf32> loc(#loc5262)
    %10604 = stablehlo.add %10600, %10603 : tensor<1x2048xf32> loc(#loc5263)
    %10605 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10606 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10607 = stablehlo.broadcast_in_dim %10606, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10608 = stablehlo.add %10605, %10607 : tensor<1xf32> loc(#loc5259)
    %10609 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10610 = stablehlo.broadcast_in_dim %10609, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10611 = stablehlo.subtract %10608, %10610 : tensor<1xf32> loc(#loc5260)
    %10612 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10613 = stablehlo.broadcast_in_dim %10612, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10614 = stablehlo.power %10613, %10611 : tensor<1xf32> loc(#loc5261)
    %10615 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10616 = stablehlo.broadcast_in_dim %10615, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10617 = stablehlo.subtract %10616, %10614 : tensor<1xf32> loc(#loc5260)
    %10618 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10619 = stablehlo.broadcast_in_dim %10618, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10620 = stablehlo.multiply %10619, %10617 : tensor<1xf32> loc(#loc5236)
    %10621 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10622 = stablehlo.broadcast_in_dim %10621, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10623 = stablehlo.power %10622, %10608 : tensor<1xf32> loc(#loc5261)
    %10624 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10625 = stablehlo.broadcast_in_dim %10624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10626 = stablehlo.subtract %10625, %10623 : tensor<1xf32> loc(#loc5260)
    %10627 = stablehlo.divide %10620, %10626 : tensor<1xf32> loc(#loc5253)
    %10628 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10629 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10630 = stablehlo.broadcast_in_dim %10629, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10631 = stablehlo.add %10628, %10630 : tensor<1xf32> loc(#loc5259)
    %10632 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10633 = stablehlo.broadcast_in_dim %10632, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10634 = stablehlo.subtract %10631, %10633 : tensor<1xf32> loc(#loc5260)
    %10635 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10636 = stablehlo.broadcast_in_dim %10635, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10637 = stablehlo.power %10636, %10634 : tensor<1xf32> loc(#loc5261)
    %10638 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10639 = stablehlo.broadcast_in_dim %10638, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10640 = stablehlo.subtract %10639, %10637 : tensor<1xf32> loc(#loc5260)
    %10641 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10642 = stablehlo.broadcast_in_dim %10641, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10643 = stablehlo.multiply %10642, %10640 : tensor<1xf32> loc(#loc5236)
    %10644 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10645 = stablehlo.broadcast_in_dim %10644, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10646 = stablehlo.power %10645, %10631 : tensor<1xf32> loc(#loc5261)
    %10647 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10648 = stablehlo.broadcast_in_dim %10647, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10649 = stablehlo.subtract %10648, %10646 : tensor<1xf32> loc(#loc5260)
    %10650 = stablehlo.divide %10643, %10649 : tensor<1xf32> loc(#loc5253)
    %10651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10652 = stablehlo.broadcast_in_dim %10651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10653 = stablehlo.subtract %10652, %10627 : tensor<1xf32> loc(#loc5260)
    %10654 = stablehlo.broadcast_in_dim %10653, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10655 = stablehlo.broadcast_in_dim %10654, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10656 = stablehlo.multiply %10655, %4612 : tensor<1x2048xf32> loc(#loc5262)
    %10657 = stablehlo.broadcast_in_dim %10627, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10658 = stablehlo.broadcast_in_dim %10657, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10659 = stablehlo.multiply %10658, %arg283 : tensor<1x2048xf32> loc(#loc5262)
    %10660 = stablehlo.add %10656, %10659 : tensor<1x2048xf32> loc(#loc5263)
    %10661 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10662 = stablehlo.broadcast_in_dim %10661, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10663 = stablehlo.subtract %10662, %10650 : tensor<1xf32> loc(#loc5260)
    %10664 = stablehlo.multiply %4612, %4612 : tensor<1x2048xf32> loc(#loc5264)
    %10665 = stablehlo.broadcast_in_dim %10663, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10666 = stablehlo.broadcast_in_dim %10665, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10667 = stablehlo.multiply %10666, %10664 : tensor<1x2048xf32> loc(#loc5262)
    %10668 = stablehlo.broadcast_in_dim %10650, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10669 = stablehlo.broadcast_in_dim %10668, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %10670 = stablehlo.multiply %10669, %arg459 : tensor<1x2048xf32> loc(#loc5262)
    %10671 = stablehlo.add %10667, %10670 : tensor<1x2048xf32> loc(#loc5263)
    %10672 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10673 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10674 = stablehlo.broadcast_in_dim %10673, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10675 = stablehlo.add %10672, %10674 : tensor<1xf32> loc(#loc5259)
    %10676 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10677 = stablehlo.broadcast_in_dim %10676, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10678 = stablehlo.subtract %10675, %10677 : tensor<1xf32> loc(#loc5260)
    %10679 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10680 = stablehlo.broadcast_in_dim %10679, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10681 = stablehlo.power %10680, %10678 : tensor<1xf32> loc(#loc5261)
    %10682 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10683 = stablehlo.broadcast_in_dim %10682, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10684 = stablehlo.subtract %10683, %10681 : tensor<1xf32> loc(#loc5260)
    %10685 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10686 = stablehlo.broadcast_in_dim %10685, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10687 = stablehlo.multiply %10686, %10684 : tensor<1xf32> loc(#loc5236)
    %10688 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10689 = stablehlo.broadcast_in_dim %10688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10690 = stablehlo.power %10689, %10675 : tensor<1xf32> loc(#loc5261)
    %10691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10692 = stablehlo.broadcast_in_dim %10691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10693 = stablehlo.subtract %10692, %10690 : tensor<1xf32> loc(#loc5260)
    %10694 = stablehlo.divide %10687, %10693 : tensor<1xf32> loc(#loc5253)
    %10695 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10696 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10697 = stablehlo.broadcast_in_dim %10696, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10698 = stablehlo.add %10695, %10697 : tensor<1xf32> loc(#loc5259)
    %10699 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10700 = stablehlo.broadcast_in_dim %10699, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10701 = stablehlo.subtract %10698, %10700 : tensor<1xf32> loc(#loc5260)
    %10702 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10703 = stablehlo.broadcast_in_dim %10702, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10704 = stablehlo.power %10703, %10701 : tensor<1xf32> loc(#loc5261)
    %10705 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10706 = stablehlo.broadcast_in_dim %10705, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10707 = stablehlo.subtract %10706, %10704 : tensor<1xf32> loc(#loc5260)
    %10708 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10709 = stablehlo.broadcast_in_dim %10708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10710 = stablehlo.multiply %10709, %10707 : tensor<1xf32> loc(#loc5236)
    %10711 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10712 = stablehlo.broadcast_in_dim %10711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10713 = stablehlo.power %10712, %10698 : tensor<1xf32> loc(#loc5261)
    %10714 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10715 = stablehlo.broadcast_in_dim %10714, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10716 = stablehlo.subtract %10715, %10713 : tensor<1xf32> loc(#loc5260)
    %10717 = stablehlo.divide %10710, %10716 : tensor<1xf32> loc(#loc5253)
    %10718 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10719 = stablehlo.broadcast_in_dim %10718, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10720 = stablehlo.subtract %10719, %10694 : tensor<1xf32> loc(#loc5260)
    %10721 = stablehlo.broadcast_in_dim %10720, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %10722 = stablehlo.broadcast_in_dim %10721, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10723 = stablehlo.multiply %10722, %4616 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10724 = stablehlo.broadcast_in_dim %10694, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %10725 = stablehlo.broadcast_in_dim %10724, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10726 = stablehlo.multiply %10725, %arg284 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10727 = stablehlo.add %10723, %10726 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %10728 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10729 = stablehlo.broadcast_in_dim %10728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10730 = stablehlo.subtract %10729, %10717 : tensor<1xf32> loc(#loc5260)
    %10731 = stablehlo.multiply %4616, %4616 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %10732 = stablehlo.broadcast_in_dim %10730, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %10733 = stablehlo.broadcast_in_dim %10732, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10734 = stablehlo.multiply %10733, %10731 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10735 = stablehlo.broadcast_in_dim %10717, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %10736 = stablehlo.broadcast_in_dim %10735, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10737 = stablehlo.multiply %10736, %arg460 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %10738 = stablehlo.add %10734, %10737 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %10739 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10740 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10741 = stablehlo.broadcast_in_dim %10740, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10742 = stablehlo.add %10739, %10741 : tensor<1xf32> loc(#loc5259)
    %10743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10744 = stablehlo.broadcast_in_dim %10743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10745 = stablehlo.subtract %10742, %10744 : tensor<1xf32> loc(#loc5260)
    %10746 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10747 = stablehlo.broadcast_in_dim %10746, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10748 = stablehlo.power %10747, %10745 : tensor<1xf32> loc(#loc5261)
    %10749 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10750 = stablehlo.broadcast_in_dim %10749, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10751 = stablehlo.subtract %10750, %10748 : tensor<1xf32> loc(#loc5260)
    %10752 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10753 = stablehlo.broadcast_in_dim %10752, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10754 = stablehlo.multiply %10753, %10751 : tensor<1xf32> loc(#loc5236)
    %10755 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10756 = stablehlo.broadcast_in_dim %10755, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10757 = stablehlo.power %10756, %10742 : tensor<1xf32> loc(#loc5261)
    %10758 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10759 = stablehlo.broadcast_in_dim %10758, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10760 = stablehlo.subtract %10759, %10757 : tensor<1xf32> loc(#loc5260)
    %10761 = stablehlo.divide %10754, %10760 : tensor<1xf32> loc(#loc5253)
    %10762 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10763 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10764 = stablehlo.broadcast_in_dim %10763, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10765 = stablehlo.add %10762, %10764 : tensor<1xf32> loc(#loc5259)
    %10766 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10767 = stablehlo.broadcast_in_dim %10766, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10768 = stablehlo.subtract %10765, %10767 : tensor<1xf32> loc(#loc5260)
    %10769 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10770 = stablehlo.broadcast_in_dim %10769, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10771 = stablehlo.power %10770, %10768 : tensor<1xf32> loc(#loc5261)
    %10772 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10773 = stablehlo.broadcast_in_dim %10772, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10774 = stablehlo.subtract %10773, %10771 : tensor<1xf32> loc(#loc5260)
    %10775 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10776 = stablehlo.broadcast_in_dim %10775, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10777 = stablehlo.multiply %10776, %10774 : tensor<1xf32> loc(#loc5236)
    %10778 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10779 = stablehlo.broadcast_in_dim %10778, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10780 = stablehlo.power %10779, %10765 : tensor<1xf32> loc(#loc5261)
    %10781 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10782 = stablehlo.broadcast_in_dim %10781, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10783 = stablehlo.subtract %10782, %10780 : tensor<1xf32> loc(#loc5260)
    %10784 = stablehlo.divide %10777, %10783 : tensor<1xf32> loc(#loc5253)
    %10785 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10786 = stablehlo.broadcast_in_dim %10785, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10787 = stablehlo.subtract %10786, %10761 : tensor<1xf32> loc(#loc5260)
    %10788 = stablehlo.broadcast_in_dim %10787, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10789 = stablehlo.broadcast_in_dim %10788, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10790 = stablehlo.multiply %10789, %4620 : tensor<1x64xf32> loc(#loc5262)
    %10791 = stablehlo.broadcast_in_dim %10761, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10792 = stablehlo.broadcast_in_dim %10791, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10793 = stablehlo.multiply %10792, %arg285 : tensor<1x64xf32> loc(#loc5262)
    %10794 = stablehlo.add %10790, %10793 : tensor<1x64xf32> loc(#loc5263)
    %10795 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10796 = stablehlo.broadcast_in_dim %10795, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10797 = stablehlo.subtract %10796, %10784 : tensor<1xf32> loc(#loc5260)
    %10798 = stablehlo.multiply %4620, %4620 : tensor<1x64xf32> loc(#loc5264)
    %10799 = stablehlo.broadcast_in_dim %10797, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10800 = stablehlo.broadcast_in_dim %10799, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10801 = stablehlo.multiply %10800, %10798 : tensor<1x64xf32> loc(#loc5262)
    %10802 = stablehlo.broadcast_in_dim %10784, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10803 = stablehlo.broadcast_in_dim %10802, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %10804 = stablehlo.multiply %10803, %arg461 : tensor<1x64xf32> loc(#loc5262)
    %10805 = stablehlo.add %10801, %10804 : tensor<1x64xf32> loc(#loc5263)
    %10806 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10807 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10808 = stablehlo.broadcast_in_dim %10807, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10809 = stablehlo.add %10806, %10808 : tensor<1xf32> loc(#loc5259)
    %10810 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10811 = stablehlo.broadcast_in_dim %10810, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10812 = stablehlo.subtract %10809, %10811 : tensor<1xf32> loc(#loc5260)
    %10813 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10814 = stablehlo.broadcast_in_dim %10813, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10815 = stablehlo.power %10814, %10812 : tensor<1xf32> loc(#loc5261)
    %10816 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10817 = stablehlo.broadcast_in_dim %10816, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10818 = stablehlo.subtract %10817, %10815 : tensor<1xf32> loc(#loc5260)
    %10819 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10820 = stablehlo.broadcast_in_dim %10819, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10821 = stablehlo.multiply %10820, %10818 : tensor<1xf32> loc(#loc5236)
    %10822 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10823 = stablehlo.broadcast_in_dim %10822, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10824 = stablehlo.power %10823, %10809 : tensor<1xf32> loc(#loc5261)
    %10825 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10826 = stablehlo.broadcast_in_dim %10825, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10827 = stablehlo.subtract %10826, %10824 : tensor<1xf32> loc(#loc5260)
    %10828 = stablehlo.divide %10821, %10827 : tensor<1xf32> loc(#loc5253)
    %10829 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10830 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10831 = stablehlo.broadcast_in_dim %10830, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10832 = stablehlo.add %10829, %10831 : tensor<1xf32> loc(#loc5259)
    %10833 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10834 = stablehlo.broadcast_in_dim %10833, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10835 = stablehlo.subtract %10832, %10834 : tensor<1xf32> loc(#loc5260)
    %10836 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10837 = stablehlo.broadcast_in_dim %10836, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10838 = stablehlo.power %10837, %10835 : tensor<1xf32> loc(#loc5261)
    %10839 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10840 = stablehlo.broadcast_in_dim %10839, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10841 = stablehlo.subtract %10840, %10838 : tensor<1xf32> loc(#loc5260)
    %10842 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10843 = stablehlo.broadcast_in_dim %10842, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10844 = stablehlo.multiply %10843, %10841 : tensor<1xf32> loc(#loc5236)
    %10845 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10846 = stablehlo.broadcast_in_dim %10845, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10847 = stablehlo.power %10846, %10832 : tensor<1xf32> loc(#loc5261)
    %10848 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10849 = stablehlo.broadcast_in_dim %10848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10850 = stablehlo.subtract %10849, %10847 : tensor<1xf32> loc(#loc5260)
    %10851 = stablehlo.divide %10844, %10850 : tensor<1xf32> loc(#loc5253)
    %10852 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10853 = stablehlo.broadcast_in_dim %10852, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10854 = stablehlo.subtract %10853, %10828 : tensor<1xf32> loc(#loc5260)
    %10855 = stablehlo.broadcast_in_dim %10854, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10856 = stablehlo.broadcast_in_dim %10855, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10857 = stablehlo.multiply %10856, %4624 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10858 = stablehlo.broadcast_in_dim %10828, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10859 = stablehlo.broadcast_in_dim %10858, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10860 = stablehlo.multiply %10859, %arg286 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10861 = stablehlo.add %10857, %10860 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %10862 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10863 = stablehlo.broadcast_in_dim %10862, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10864 = stablehlo.subtract %10863, %10851 : tensor<1xf32> loc(#loc5260)
    %10865 = stablehlo.multiply %4624, %4624 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %10866 = stablehlo.broadcast_in_dim %10864, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10867 = stablehlo.broadcast_in_dim %10866, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10868 = stablehlo.multiply %10867, %10865 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10869 = stablehlo.broadcast_in_dim %10851, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %10870 = stablehlo.broadcast_in_dim %10869, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %10871 = stablehlo.multiply %10870, %arg462 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %10872 = stablehlo.add %10868, %10871 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %10873 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10874 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10875 = stablehlo.broadcast_in_dim %10874, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10876 = stablehlo.add %10873, %10875 : tensor<1xf32> loc(#loc5259)
    %10877 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10878 = stablehlo.broadcast_in_dim %10877, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10879 = stablehlo.subtract %10876, %10878 : tensor<1xf32> loc(#loc5260)
    %10880 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10881 = stablehlo.broadcast_in_dim %10880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10882 = stablehlo.power %10881, %10879 : tensor<1xf32> loc(#loc5261)
    %10883 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10884 = stablehlo.broadcast_in_dim %10883, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10885 = stablehlo.subtract %10884, %10882 : tensor<1xf32> loc(#loc5260)
    %10886 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10887 = stablehlo.broadcast_in_dim %10886, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10888 = stablehlo.multiply %10887, %10885 : tensor<1xf32> loc(#loc5236)
    %10889 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10890 = stablehlo.broadcast_in_dim %10889, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10891 = stablehlo.power %10890, %10876 : tensor<1xf32> loc(#loc5261)
    %10892 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10893 = stablehlo.broadcast_in_dim %10892, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10894 = stablehlo.subtract %10893, %10891 : tensor<1xf32> loc(#loc5260)
    %10895 = stablehlo.divide %10888, %10894 : tensor<1xf32> loc(#loc5253)
    %10896 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10897 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10898 = stablehlo.broadcast_in_dim %10897, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10899 = stablehlo.add %10896, %10898 : tensor<1xf32> loc(#loc5259)
    %10900 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10901 = stablehlo.broadcast_in_dim %10900, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10902 = stablehlo.subtract %10899, %10901 : tensor<1xf32> loc(#loc5260)
    %10903 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10904 = stablehlo.broadcast_in_dim %10903, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10905 = stablehlo.power %10904, %10902 : tensor<1xf32> loc(#loc5261)
    %10906 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10907 = stablehlo.broadcast_in_dim %10906, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10908 = stablehlo.subtract %10907, %10905 : tensor<1xf32> loc(#loc5260)
    %10909 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10910 = stablehlo.broadcast_in_dim %10909, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10911 = stablehlo.multiply %10910, %10908 : tensor<1xf32> loc(#loc5236)
    %10912 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10913 = stablehlo.broadcast_in_dim %10912, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10914 = stablehlo.power %10913, %10899 : tensor<1xf32> loc(#loc5261)
    %10915 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10916 = stablehlo.broadcast_in_dim %10915, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10917 = stablehlo.subtract %10916, %10914 : tensor<1xf32> loc(#loc5260)
    %10918 = stablehlo.divide %10911, %10917 : tensor<1xf32> loc(#loc5253)
    %10919 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10920 = stablehlo.broadcast_in_dim %10919, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10921 = stablehlo.subtract %10920, %10895 : tensor<1xf32> loc(#loc5260)
    %10922 = stablehlo.broadcast_in_dim %10921, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10923 = stablehlo.broadcast_in_dim %10922, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10924 = stablehlo.multiply %10923, %4628 : tensor<1x8192xf32> loc(#loc5262)
    %10925 = stablehlo.broadcast_in_dim %10895, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10926 = stablehlo.broadcast_in_dim %10925, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10927 = stablehlo.multiply %10926, %arg287 : tensor<1x8192xf32> loc(#loc5262)
    %10928 = stablehlo.add %10924, %10927 : tensor<1x8192xf32> loc(#loc5263)
    %10929 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10930 = stablehlo.broadcast_in_dim %10929, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10931 = stablehlo.subtract %10930, %10918 : tensor<1xf32> loc(#loc5260)
    %10932 = stablehlo.multiply %4628, %4628 : tensor<1x8192xf32> loc(#loc5264)
    %10933 = stablehlo.broadcast_in_dim %10931, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10934 = stablehlo.broadcast_in_dim %10933, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10935 = stablehlo.multiply %10934, %10932 : tensor<1x8192xf32> loc(#loc5262)
    %10936 = stablehlo.broadcast_in_dim %10918, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %10937 = stablehlo.broadcast_in_dim %10936, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %10938 = stablehlo.multiply %10937, %arg463 : tensor<1x8192xf32> loc(#loc5262)
    %10939 = stablehlo.add %10935, %10938 : tensor<1x8192xf32> loc(#loc5263)
    %10940 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10941 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10942 = stablehlo.broadcast_in_dim %10941, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10943 = stablehlo.add %10940, %10942 : tensor<1xf32> loc(#loc5259)
    %10944 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10945 = stablehlo.broadcast_in_dim %10944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10946 = stablehlo.subtract %10943, %10945 : tensor<1xf32> loc(#loc5260)
    %10947 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10948 = stablehlo.broadcast_in_dim %10947, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10949 = stablehlo.power %10948, %10946 : tensor<1xf32> loc(#loc5261)
    %10950 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10951 = stablehlo.broadcast_in_dim %10950, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10952 = stablehlo.subtract %10951, %10949 : tensor<1xf32> loc(#loc5260)
    %10953 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10954 = stablehlo.broadcast_in_dim %10953, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10955 = stablehlo.multiply %10954, %10952 : tensor<1xf32> loc(#loc5236)
    %10956 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %10957 = stablehlo.broadcast_in_dim %10956, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10958 = stablehlo.power %10957, %10943 : tensor<1xf32> loc(#loc5261)
    %10959 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10960 = stablehlo.broadcast_in_dim %10959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10961 = stablehlo.subtract %10960, %10958 : tensor<1xf32> loc(#loc5260)
    %10962 = stablehlo.divide %10955, %10961 : tensor<1xf32> loc(#loc5253)
    %10963 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %10964 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10965 = stablehlo.broadcast_in_dim %10964, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %10966 = stablehlo.add %10963, %10965 : tensor<1xf32> loc(#loc5259)
    %10967 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10968 = stablehlo.broadcast_in_dim %10967, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10969 = stablehlo.subtract %10966, %10968 : tensor<1xf32> loc(#loc5260)
    %10970 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10971 = stablehlo.broadcast_in_dim %10970, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10972 = stablehlo.power %10971, %10969 : tensor<1xf32> loc(#loc5261)
    %10973 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10974 = stablehlo.broadcast_in_dim %10973, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10975 = stablehlo.subtract %10974, %10972 : tensor<1xf32> loc(#loc5260)
    %10976 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10977 = stablehlo.broadcast_in_dim %10976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %10978 = stablehlo.multiply %10977, %10975 : tensor<1xf32> loc(#loc5236)
    %10979 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %10980 = stablehlo.broadcast_in_dim %10979, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %10981 = stablehlo.power %10980, %10966 : tensor<1xf32> loc(#loc5261)
    %10982 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10983 = stablehlo.broadcast_in_dim %10982, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10984 = stablehlo.subtract %10983, %10981 : tensor<1xf32> loc(#loc5260)
    %10985 = stablehlo.divide %10978, %10984 : tensor<1xf32> loc(#loc5253)
    %10986 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10987 = stablehlo.broadcast_in_dim %10986, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10988 = stablehlo.subtract %10987, %10962 : tensor<1xf32> loc(#loc5260)
    %10989 = stablehlo.broadcast_in_dim %10988, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10990 = stablehlo.broadcast_in_dim %10989, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %10991 = stablehlo.multiply %10990, %4632 : tensor<1x2048x8192xf32> loc(#loc5262)
    %10992 = stablehlo.broadcast_in_dim %10962, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %10993 = stablehlo.broadcast_in_dim %10992, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %10994 = stablehlo.multiply %10993, %arg288 : tensor<1x2048x8192xf32> loc(#loc5262)
    %10995 = stablehlo.add %10991, %10994 : tensor<1x2048x8192xf32> loc(#loc5263)
    %10996 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %10997 = stablehlo.broadcast_in_dim %10996, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %10998 = stablehlo.subtract %10997, %10985 : tensor<1xf32> loc(#loc5260)
    %10999 = stablehlo.multiply %4632, %4632 : tensor<1x2048x8192xf32> loc(#loc5264)
    %11000 = stablehlo.broadcast_in_dim %10998, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11001 = stablehlo.broadcast_in_dim %11000, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %11002 = stablehlo.multiply %11001, %10999 : tensor<1x2048x8192xf32> loc(#loc5262)
    %11003 = stablehlo.broadcast_in_dim %10985, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11004 = stablehlo.broadcast_in_dim %11003, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %11005 = stablehlo.multiply %11004, %arg464 : tensor<1x2048x8192xf32> loc(#loc5262)
    %11006 = stablehlo.add %11002, %11005 : tensor<1x2048x8192xf32> loc(#loc5263)
    %11007 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11008 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11009 = stablehlo.broadcast_in_dim %11008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11010 = stablehlo.add %11007, %11009 : tensor<1xf32> loc(#loc5259)
    %11011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11012 = stablehlo.broadcast_in_dim %11011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11013 = stablehlo.subtract %11010, %11012 : tensor<1xf32> loc(#loc5260)
    %11014 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11015 = stablehlo.broadcast_in_dim %11014, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11016 = stablehlo.power %11015, %11013 : tensor<1xf32> loc(#loc5261)
    %11017 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11018 = stablehlo.broadcast_in_dim %11017, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11019 = stablehlo.subtract %11018, %11016 : tensor<1xf32> loc(#loc5260)
    %11020 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11021 = stablehlo.broadcast_in_dim %11020, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11022 = stablehlo.multiply %11021, %11019 : tensor<1xf32> loc(#loc5236)
    %11023 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11024 = stablehlo.broadcast_in_dim %11023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11025 = stablehlo.power %11024, %11010 : tensor<1xf32> loc(#loc5261)
    %11026 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11027 = stablehlo.broadcast_in_dim %11026, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11028 = stablehlo.subtract %11027, %11025 : tensor<1xf32> loc(#loc5260)
    %11029 = stablehlo.divide %11022, %11028 : tensor<1xf32> loc(#loc5253)
    %11030 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11032 = stablehlo.broadcast_in_dim %11031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11033 = stablehlo.add %11030, %11032 : tensor<1xf32> loc(#loc5259)
    %11034 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11035 = stablehlo.broadcast_in_dim %11034, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11036 = stablehlo.subtract %11033, %11035 : tensor<1xf32> loc(#loc5260)
    %11037 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11038 = stablehlo.broadcast_in_dim %11037, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11039 = stablehlo.power %11038, %11036 : tensor<1xf32> loc(#loc5261)
    %11040 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11041 = stablehlo.broadcast_in_dim %11040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11042 = stablehlo.subtract %11041, %11039 : tensor<1xf32> loc(#loc5260)
    %11043 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11044 = stablehlo.broadcast_in_dim %11043, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11045 = stablehlo.multiply %11044, %11042 : tensor<1xf32> loc(#loc5236)
    %11046 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11047 = stablehlo.broadcast_in_dim %11046, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11048 = stablehlo.power %11047, %11033 : tensor<1xf32> loc(#loc5261)
    %11049 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11050 = stablehlo.broadcast_in_dim %11049, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11051 = stablehlo.subtract %11050, %11048 : tensor<1xf32> loc(#loc5260)
    %11052 = stablehlo.divide %11045, %11051 : tensor<1xf32> loc(#loc5253)
    %11053 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11054 = stablehlo.broadcast_in_dim %11053, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11055 = stablehlo.subtract %11054, %11029 : tensor<1xf32> loc(#loc5260)
    %11056 = stablehlo.broadcast_in_dim %11055, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11057 = stablehlo.broadcast_in_dim %11056, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11058 = stablehlo.multiply %11057, %4636 : tensor<1x2048xf32> loc(#loc5262)
    %11059 = stablehlo.broadcast_in_dim %11029, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11060 = stablehlo.broadcast_in_dim %11059, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11061 = stablehlo.multiply %11060, %arg289 : tensor<1x2048xf32> loc(#loc5262)
    %11062 = stablehlo.add %11058, %11061 : tensor<1x2048xf32> loc(#loc5263)
    %11063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11064 = stablehlo.broadcast_in_dim %11063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11065 = stablehlo.subtract %11064, %11052 : tensor<1xf32> loc(#loc5260)
    %11066 = stablehlo.multiply %4636, %4636 : tensor<1x2048xf32> loc(#loc5264)
    %11067 = stablehlo.broadcast_in_dim %11065, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11068 = stablehlo.broadcast_in_dim %11067, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11069 = stablehlo.multiply %11068, %11066 : tensor<1x2048xf32> loc(#loc5262)
    %11070 = stablehlo.broadcast_in_dim %11052, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11071 = stablehlo.broadcast_in_dim %11070, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11072 = stablehlo.multiply %11071, %arg465 : tensor<1x2048xf32> loc(#loc5262)
    %11073 = stablehlo.add %11069, %11072 : tensor<1x2048xf32> loc(#loc5263)
    %11074 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11075 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11076 = stablehlo.broadcast_in_dim %11075, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11077 = stablehlo.add %11074, %11076 : tensor<1xf32> loc(#loc5259)
    %11078 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11079 = stablehlo.broadcast_in_dim %11078, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11080 = stablehlo.subtract %11077, %11079 : tensor<1xf32> loc(#loc5260)
    %11081 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11082 = stablehlo.broadcast_in_dim %11081, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11083 = stablehlo.power %11082, %11080 : tensor<1xf32> loc(#loc5261)
    %11084 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11085 = stablehlo.broadcast_in_dim %11084, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11086 = stablehlo.subtract %11085, %11083 : tensor<1xf32> loc(#loc5260)
    %11087 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11088 = stablehlo.broadcast_in_dim %11087, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11089 = stablehlo.multiply %11088, %11086 : tensor<1xf32> loc(#loc5236)
    %11090 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11091 = stablehlo.broadcast_in_dim %11090, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11092 = stablehlo.power %11091, %11077 : tensor<1xf32> loc(#loc5261)
    %11093 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11094 = stablehlo.broadcast_in_dim %11093, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11095 = stablehlo.subtract %11094, %11092 : tensor<1xf32> loc(#loc5260)
    %11096 = stablehlo.divide %11089, %11095 : tensor<1xf32> loc(#loc5253)
    %11097 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11098 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11099 = stablehlo.broadcast_in_dim %11098, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11100 = stablehlo.add %11097, %11099 : tensor<1xf32> loc(#loc5259)
    %11101 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11102 = stablehlo.broadcast_in_dim %11101, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11103 = stablehlo.subtract %11100, %11102 : tensor<1xf32> loc(#loc5260)
    %11104 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11105 = stablehlo.broadcast_in_dim %11104, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11106 = stablehlo.power %11105, %11103 : tensor<1xf32> loc(#loc5261)
    %11107 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11108 = stablehlo.broadcast_in_dim %11107, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11109 = stablehlo.subtract %11108, %11106 : tensor<1xf32> loc(#loc5260)
    %11110 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11111 = stablehlo.broadcast_in_dim %11110, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11112 = stablehlo.multiply %11111, %11109 : tensor<1xf32> loc(#loc5236)
    %11113 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11114 = stablehlo.broadcast_in_dim %11113, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11115 = stablehlo.power %11114, %11100 : tensor<1xf32> loc(#loc5261)
    %11116 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11117 = stablehlo.broadcast_in_dim %11116, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11118 = stablehlo.subtract %11117, %11115 : tensor<1xf32> loc(#loc5260)
    %11119 = stablehlo.divide %11112, %11118 : tensor<1xf32> loc(#loc5253)
    %11120 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11121 = stablehlo.broadcast_in_dim %11120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11122 = stablehlo.subtract %11121, %11096 : tensor<1xf32> loc(#loc5260)
    %11123 = stablehlo.broadcast_in_dim %11122, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11124 = stablehlo.broadcast_in_dim %11123, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11125 = stablehlo.multiply %11124, %4640 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11126 = stablehlo.broadcast_in_dim %11096, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11127 = stablehlo.broadcast_in_dim %11126, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11128 = stablehlo.multiply %11127, %arg290 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11129 = stablehlo.add %11125, %11128 : tensor<1x8192x2048xf32> loc(#loc5263)
    %11130 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11131 = stablehlo.broadcast_in_dim %11130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11132 = stablehlo.subtract %11131, %11119 : tensor<1xf32> loc(#loc5260)
    %11133 = stablehlo.multiply %4640, %4640 : tensor<1x8192x2048xf32> loc(#loc5264)
    %11134 = stablehlo.broadcast_in_dim %11132, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11135 = stablehlo.broadcast_in_dim %11134, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11136 = stablehlo.multiply %11135, %11133 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11137 = stablehlo.broadcast_in_dim %11119, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11138 = stablehlo.broadcast_in_dim %11137, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11139 = stablehlo.multiply %11138, %arg466 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11140 = stablehlo.add %11136, %11139 : tensor<1x8192x2048xf32> loc(#loc5263)
    %11141 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11142 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11143 = stablehlo.broadcast_in_dim %11142, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11144 = stablehlo.add %11141, %11143 : tensor<1xf32> loc(#loc5259)
    %11145 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11146 = stablehlo.broadcast_in_dim %11145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11147 = stablehlo.subtract %11144, %11146 : tensor<1xf32> loc(#loc5260)
    %11148 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11149 = stablehlo.broadcast_in_dim %11148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11150 = stablehlo.power %11149, %11147 : tensor<1xf32> loc(#loc5261)
    %11151 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11152 = stablehlo.broadcast_in_dim %11151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11153 = stablehlo.subtract %11152, %11150 : tensor<1xf32> loc(#loc5260)
    %11154 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11155 = stablehlo.broadcast_in_dim %11154, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11156 = stablehlo.multiply %11155, %11153 : tensor<1xf32> loc(#loc5236)
    %11157 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11158 = stablehlo.broadcast_in_dim %11157, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11159 = stablehlo.power %11158, %11144 : tensor<1xf32> loc(#loc5261)
    %11160 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11161 = stablehlo.broadcast_in_dim %11160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11162 = stablehlo.subtract %11161, %11159 : tensor<1xf32> loc(#loc5260)
    %11163 = stablehlo.divide %11156, %11162 : tensor<1xf32> loc(#loc5253)
    %11164 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11165 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11166 = stablehlo.broadcast_in_dim %11165, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11167 = stablehlo.add %11164, %11166 : tensor<1xf32> loc(#loc5259)
    %11168 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11169 = stablehlo.broadcast_in_dim %11168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11170 = stablehlo.subtract %11167, %11169 : tensor<1xf32> loc(#loc5260)
    %11171 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11172 = stablehlo.broadcast_in_dim %11171, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11173 = stablehlo.power %11172, %11170 : tensor<1xf32> loc(#loc5261)
    %11174 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11175 = stablehlo.broadcast_in_dim %11174, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11176 = stablehlo.subtract %11175, %11173 : tensor<1xf32> loc(#loc5260)
    %11177 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11178 = stablehlo.broadcast_in_dim %11177, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11179 = stablehlo.multiply %11178, %11176 : tensor<1xf32> loc(#loc5236)
    %11180 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11181 = stablehlo.broadcast_in_dim %11180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11182 = stablehlo.power %11181, %11167 : tensor<1xf32> loc(#loc5261)
    %11183 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11184 = stablehlo.broadcast_in_dim %11183, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11185 = stablehlo.subtract %11184, %11182 : tensor<1xf32> loc(#loc5260)
    %11186 = stablehlo.divide %11179, %11185 : tensor<1xf32> loc(#loc5253)
    %11187 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11188 = stablehlo.broadcast_in_dim %11187, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11189 = stablehlo.subtract %11188, %11163 : tensor<1xf32> loc(#loc5260)
    %11190 = stablehlo.broadcast_in_dim %11189, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11191 = stablehlo.broadcast_in_dim %11190, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11192 = stablehlo.multiply %11191, %4644 : tensor<1x2048xf32> loc(#loc5262)
    %11193 = stablehlo.broadcast_in_dim %11163, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11194 = stablehlo.broadcast_in_dim %11193, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11195 = stablehlo.multiply %11194, %arg291 : tensor<1x2048xf32> loc(#loc5262)
    %11196 = stablehlo.add %11192, %11195 : tensor<1x2048xf32> loc(#loc5263)
    %11197 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11198 = stablehlo.broadcast_in_dim %11197, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11199 = stablehlo.subtract %11198, %11186 : tensor<1xf32> loc(#loc5260)
    %11200 = stablehlo.multiply %4644, %4644 : tensor<1x2048xf32> loc(#loc5264)
    %11201 = stablehlo.broadcast_in_dim %11199, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11202 = stablehlo.broadcast_in_dim %11201, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11203 = stablehlo.multiply %11202, %11200 : tensor<1x2048xf32> loc(#loc5262)
    %11204 = stablehlo.broadcast_in_dim %11186, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11205 = stablehlo.broadcast_in_dim %11204, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11206 = stablehlo.multiply %11205, %arg467 : tensor<1x2048xf32> loc(#loc5262)
    %11207 = stablehlo.add %11203, %11206 : tensor<1x2048xf32> loc(#loc5263)
    %11208 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11209 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11210 = stablehlo.broadcast_in_dim %11209, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11211 = stablehlo.add %11208, %11210 : tensor<1xf32> loc(#loc5259)
    %11212 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11213 = stablehlo.broadcast_in_dim %11212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11214 = stablehlo.subtract %11211, %11213 : tensor<1xf32> loc(#loc5260)
    %11215 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11216 = stablehlo.broadcast_in_dim %11215, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11217 = stablehlo.power %11216, %11214 : tensor<1xf32> loc(#loc5261)
    %11218 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11219 = stablehlo.broadcast_in_dim %11218, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11220 = stablehlo.subtract %11219, %11217 : tensor<1xf32> loc(#loc5260)
    %11221 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11222 = stablehlo.broadcast_in_dim %11221, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11223 = stablehlo.multiply %11222, %11220 : tensor<1xf32> loc(#loc5236)
    %11224 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11225 = stablehlo.broadcast_in_dim %11224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11226 = stablehlo.power %11225, %11211 : tensor<1xf32> loc(#loc5261)
    %11227 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11228 = stablehlo.broadcast_in_dim %11227, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11229 = stablehlo.subtract %11228, %11226 : tensor<1xf32> loc(#loc5260)
    %11230 = stablehlo.divide %11223, %11229 : tensor<1xf32> loc(#loc5253)
    %11231 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11232 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11233 = stablehlo.broadcast_in_dim %11232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11234 = stablehlo.add %11231, %11233 : tensor<1xf32> loc(#loc5259)
    %11235 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11236 = stablehlo.broadcast_in_dim %11235, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11237 = stablehlo.subtract %11234, %11236 : tensor<1xf32> loc(#loc5260)
    %11238 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11239 = stablehlo.broadcast_in_dim %11238, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11240 = stablehlo.power %11239, %11237 : tensor<1xf32> loc(#loc5261)
    %11241 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11242 = stablehlo.broadcast_in_dim %11241, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11243 = stablehlo.subtract %11242, %11240 : tensor<1xf32> loc(#loc5260)
    %11244 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11245 = stablehlo.broadcast_in_dim %11244, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11246 = stablehlo.multiply %11245, %11243 : tensor<1xf32> loc(#loc5236)
    %11247 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11248 = stablehlo.broadcast_in_dim %11247, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11249 = stablehlo.power %11248, %11234 : tensor<1xf32> loc(#loc5261)
    %11250 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11251 = stablehlo.broadcast_in_dim %11250, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11252 = stablehlo.subtract %11251, %11249 : tensor<1xf32> loc(#loc5260)
    %11253 = stablehlo.divide %11246, %11252 : tensor<1xf32> loc(#loc5253)
    %11254 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11255 = stablehlo.broadcast_in_dim %11254, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11256 = stablehlo.subtract %11255, %11230 : tensor<1xf32> loc(#loc5260)
    %11257 = stablehlo.broadcast_in_dim %11256, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11258 = stablehlo.broadcast_in_dim %11257, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11259 = stablehlo.multiply %11258, %4648 : tensor<1x2048xf32> loc(#loc5262)
    %11260 = stablehlo.broadcast_in_dim %11230, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11261 = stablehlo.broadcast_in_dim %11260, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11262 = stablehlo.multiply %11261, %arg292 : tensor<1x2048xf32> loc(#loc5262)
    %11263 = stablehlo.add %11259, %11262 : tensor<1x2048xf32> loc(#loc5263)
    %11264 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11265 = stablehlo.broadcast_in_dim %11264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11266 = stablehlo.subtract %11265, %11253 : tensor<1xf32> loc(#loc5260)
    %11267 = stablehlo.multiply %4648, %4648 : tensor<1x2048xf32> loc(#loc5264)
    %11268 = stablehlo.broadcast_in_dim %11266, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11269 = stablehlo.broadcast_in_dim %11268, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11270 = stablehlo.multiply %11269, %11267 : tensor<1x2048xf32> loc(#loc5262)
    %11271 = stablehlo.broadcast_in_dim %11253, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11272 = stablehlo.broadcast_in_dim %11271, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11273 = stablehlo.multiply %11272, %arg468 : tensor<1x2048xf32> loc(#loc5262)
    %11274 = stablehlo.add %11270, %11273 : tensor<1x2048xf32> loc(#loc5263)
    %11275 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11276 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11277 = stablehlo.broadcast_in_dim %11276, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11278 = stablehlo.add %11275, %11277 : tensor<1xf32> loc(#loc5259)
    %11279 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11280 = stablehlo.broadcast_in_dim %11279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11281 = stablehlo.subtract %11278, %11280 : tensor<1xf32> loc(#loc5260)
    %11282 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11283 = stablehlo.broadcast_in_dim %11282, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11284 = stablehlo.power %11283, %11281 : tensor<1xf32> loc(#loc5261)
    %11285 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11286 = stablehlo.broadcast_in_dim %11285, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11287 = stablehlo.subtract %11286, %11284 : tensor<1xf32> loc(#loc5260)
    %11288 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11289 = stablehlo.broadcast_in_dim %11288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11290 = stablehlo.multiply %11289, %11287 : tensor<1xf32> loc(#loc5236)
    %11291 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11292 = stablehlo.broadcast_in_dim %11291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11293 = stablehlo.power %11292, %11278 : tensor<1xf32> loc(#loc5261)
    %11294 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11295 = stablehlo.broadcast_in_dim %11294, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11296 = stablehlo.subtract %11295, %11293 : tensor<1xf32> loc(#loc5260)
    %11297 = stablehlo.divide %11290, %11296 : tensor<1xf32> loc(#loc5253)
    %11298 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11299 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11300 = stablehlo.broadcast_in_dim %11299, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11301 = stablehlo.add %11298, %11300 : tensor<1xf32> loc(#loc5259)
    %11302 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11303 = stablehlo.broadcast_in_dim %11302, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11304 = stablehlo.subtract %11301, %11303 : tensor<1xf32> loc(#loc5260)
    %11305 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11306 = stablehlo.broadcast_in_dim %11305, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11307 = stablehlo.power %11306, %11304 : tensor<1xf32> loc(#loc5261)
    %11308 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11309 = stablehlo.broadcast_in_dim %11308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11310 = stablehlo.subtract %11309, %11307 : tensor<1xf32> loc(#loc5260)
    %11311 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11312 = stablehlo.broadcast_in_dim %11311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11313 = stablehlo.multiply %11312, %11310 : tensor<1xf32> loc(#loc5236)
    %11314 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11315 = stablehlo.broadcast_in_dim %11314, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11316 = stablehlo.power %11315, %11301 : tensor<1xf32> loc(#loc5261)
    %11317 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11318 = stablehlo.broadcast_in_dim %11317, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11319 = stablehlo.subtract %11318, %11316 : tensor<1xf32> loc(#loc5260)
    %11320 = stablehlo.divide %11313, %11319 : tensor<1xf32> loc(#loc5253)
    %11321 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11322 = stablehlo.broadcast_in_dim %11321, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11323 = stablehlo.subtract %11322, %11297 : tensor<1xf32> loc(#loc5260)
    %11324 = stablehlo.broadcast_in_dim %11323, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11325 = stablehlo.broadcast_in_dim %11324, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11326 = stablehlo.multiply %11325, %4652 : tensor<1x2048xf32> loc(#loc5262)
    %11327 = stablehlo.broadcast_in_dim %11297, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11328 = stablehlo.broadcast_in_dim %11327, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11329 = stablehlo.multiply %11328, %arg293 : tensor<1x2048xf32> loc(#loc5262)
    %11330 = stablehlo.add %11326, %11329 : tensor<1x2048xf32> loc(#loc5263)
    %11331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11332 = stablehlo.broadcast_in_dim %11331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11333 = stablehlo.subtract %11332, %11320 : tensor<1xf32> loc(#loc5260)
    %11334 = stablehlo.multiply %4652, %4652 : tensor<1x2048xf32> loc(#loc5264)
    %11335 = stablehlo.broadcast_in_dim %11333, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11336 = stablehlo.broadcast_in_dim %11335, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11337 = stablehlo.multiply %11336, %11334 : tensor<1x2048xf32> loc(#loc5262)
    %11338 = stablehlo.broadcast_in_dim %11320, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11339 = stablehlo.broadcast_in_dim %11338, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11340 = stablehlo.multiply %11339, %arg469 : tensor<1x2048xf32> loc(#loc5262)
    %11341 = stablehlo.add %11337, %11340 : tensor<1x2048xf32> loc(#loc5263)
    %11342 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11343 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11344 = stablehlo.broadcast_in_dim %11343, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11345 = stablehlo.add %11342, %11344 : tensor<1xf32> loc(#loc5259)
    %11346 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11347 = stablehlo.broadcast_in_dim %11346, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11348 = stablehlo.subtract %11345, %11347 : tensor<1xf32> loc(#loc5260)
    %11349 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11350 = stablehlo.broadcast_in_dim %11349, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11351 = stablehlo.power %11350, %11348 : tensor<1xf32> loc(#loc5261)
    %11352 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11353 = stablehlo.broadcast_in_dim %11352, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11354 = stablehlo.subtract %11353, %11351 : tensor<1xf32> loc(#loc5260)
    %11355 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11356 = stablehlo.broadcast_in_dim %11355, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11357 = stablehlo.multiply %11356, %11354 : tensor<1xf32> loc(#loc5236)
    %11358 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11359 = stablehlo.broadcast_in_dim %11358, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11360 = stablehlo.power %11359, %11345 : tensor<1xf32> loc(#loc5261)
    %11361 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11362 = stablehlo.broadcast_in_dim %11361, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11363 = stablehlo.subtract %11362, %11360 : tensor<1xf32> loc(#loc5260)
    %11364 = stablehlo.divide %11357, %11363 : tensor<1xf32> loc(#loc5253)
    %11365 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11366 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11367 = stablehlo.broadcast_in_dim %11366, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11368 = stablehlo.add %11365, %11367 : tensor<1xf32> loc(#loc5259)
    %11369 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11370 = stablehlo.broadcast_in_dim %11369, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11371 = stablehlo.subtract %11368, %11370 : tensor<1xf32> loc(#loc5260)
    %11372 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11373 = stablehlo.broadcast_in_dim %11372, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11374 = stablehlo.power %11373, %11371 : tensor<1xf32> loc(#loc5261)
    %11375 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11376 = stablehlo.broadcast_in_dim %11375, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11377 = stablehlo.subtract %11376, %11374 : tensor<1xf32> loc(#loc5260)
    %11378 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11379 = stablehlo.broadcast_in_dim %11378, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11380 = stablehlo.multiply %11379, %11377 : tensor<1xf32> loc(#loc5236)
    %11381 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11382 = stablehlo.broadcast_in_dim %11381, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11383 = stablehlo.power %11382, %11368 : tensor<1xf32> loc(#loc5261)
    %11384 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11385 = stablehlo.broadcast_in_dim %11384, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11386 = stablehlo.subtract %11385, %11383 : tensor<1xf32> loc(#loc5260)
    %11387 = stablehlo.divide %11380, %11386 : tensor<1xf32> loc(#loc5253)
    %11388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11389 = stablehlo.broadcast_in_dim %11388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11390 = stablehlo.subtract %11389, %11364 : tensor<1xf32> loc(#loc5260)
    %11391 = stablehlo.broadcast_in_dim %11390, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11392 = stablehlo.broadcast_in_dim %11391, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11393 = stablehlo.multiply %11392, %4656 : tensor<1x2048xf32> loc(#loc5262)
    %11394 = stablehlo.broadcast_in_dim %11364, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11395 = stablehlo.broadcast_in_dim %11394, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11396 = stablehlo.multiply %11395, %arg294 : tensor<1x2048xf32> loc(#loc5262)
    %11397 = stablehlo.add %11393, %11396 : tensor<1x2048xf32> loc(#loc5263)
    %11398 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11399 = stablehlo.broadcast_in_dim %11398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11400 = stablehlo.subtract %11399, %11387 : tensor<1xf32> loc(#loc5260)
    %11401 = stablehlo.multiply %4656, %4656 : tensor<1x2048xf32> loc(#loc5264)
    %11402 = stablehlo.broadcast_in_dim %11400, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11403 = stablehlo.broadcast_in_dim %11402, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11404 = stablehlo.multiply %11403, %11401 : tensor<1x2048xf32> loc(#loc5262)
    %11405 = stablehlo.broadcast_in_dim %11387, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11406 = stablehlo.broadcast_in_dim %11405, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11407 = stablehlo.multiply %11406, %arg470 : tensor<1x2048xf32> loc(#loc5262)
    %11408 = stablehlo.add %11404, %11407 : tensor<1x2048xf32> loc(#loc5263)
    %11409 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11410 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11411 = stablehlo.broadcast_in_dim %11410, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11412 = stablehlo.add %11409, %11411 : tensor<1xf32> loc(#loc5259)
    %11413 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11414 = stablehlo.broadcast_in_dim %11413, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11415 = stablehlo.subtract %11412, %11414 : tensor<1xf32> loc(#loc5260)
    %11416 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11417 = stablehlo.broadcast_in_dim %11416, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11418 = stablehlo.power %11417, %11415 : tensor<1xf32> loc(#loc5261)
    %11419 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11420 = stablehlo.broadcast_in_dim %11419, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11421 = stablehlo.subtract %11420, %11418 : tensor<1xf32> loc(#loc5260)
    %11422 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11423 = stablehlo.broadcast_in_dim %11422, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11424 = stablehlo.multiply %11423, %11421 : tensor<1xf32> loc(#loc5236)
    %11425 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11426 = stablehlo.broadcast_in_dim %11425, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11427 = stablehlo.power %11426, %11412 : tensor<1xf32> loc(#loc5261)
    %11428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11429 = stablehlo.broadcast_in_dim %11428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11430 = stablehlo.subtract %11429, %11427 : tensor<1xf32> loc(#loc5260)
    %11431 = stablehlo.divide %11424, %11430 : tensor<1xf32> loc(#loc5253)
    %11432 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11433 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11434 = stablehlo.broadcast_in_dim %11433, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11435 = stablehlo.add %11432, %11434 : tensor<1xf32> loc(#loc5259)
    %11436 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11437 = stablehlo.broadcast_in_dim %11436, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11438 = stablehlo.subtract %11435, %11437 : tensor<1xf32> loc(#loc5260)
    %11439 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11440 = stablehlo.broadcast_in_dim %11439, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11441 = stablehlo.power %11440, %11438 : tensor<1xf32> loc(#loc5261)
    %11442 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11443 = stablehlo.broadcast_in_dim %11442, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11444 = stablehlo.subtract %11443, %11441 : tensor<1xf32> loc(#loc5260)
    %11445 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11446 = stablehlo.broadcast_in_dim %11445, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11447 = stablehlo.multiply %11446, %11444 : tensor<1xf32> loc(#loc5236)
    %11448 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11449 = stablehlo.broadcast_in_dim %11448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11450 = stablehlo.power %11449, %11435 : tensor<1xf32> loc(#loc5261)
    %11451 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11452 = stablehlo.broadcast_in_dim %11451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11453 = stablehlo.subtract %11452, %11450 : tensor<1xf32> loc(#loc5260)
    %11454 = stablehlo.divide %11447, %11453 : tensor<1xf32> loc(#loc5253)
    %11455 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11456 = stablehlo.broadcast_in_dim %11455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11457 = stablehlo.subtract %11456, %11431 : tensor<1xf32> loc(#loc5260)
    %11458 = stablehlo.broadcast_in_dim %11457, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %11459 = stablehlo.broadcast_in_dim %11458, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11460 = stablehlo.multiply %11459, %4660 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11461 = stablehlo.broadcast_in_dim %11431, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %11462 = stablehlo.broadcast_in_dim %11461, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11463 = stablehlo.multiply %11462, %arg295 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11464 = stablehlo.add %11460, %11463 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %11465 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11466 = stablehlo.broadcast_in_dim %11465, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11467 = stablehlo.subtract %11466, %11454 : tensor<1xf32> loc(#loc5260)
    %11468 = stablehlo.multiply %4660, %4660 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %11469 = stablehlo.broadcast_in_dim %11467, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %11470 = stablehlo.broadcast_in_dim %11469, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11471 = stablehlo.multiply %11470, %11468 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11472 = stablehlo.broadcast_in_dim %11454, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %11473 = stablehlo.broadcast_in_dim %11472, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11474 = stablehlo.multiply %11473, %arg471 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %11475 = stablehlo.add %11471, %11474 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %11476 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11477 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11478 = stablehlo.broadcast_in_dim %11477, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11479 = stablehlo.add %11476, %11478 : tensor<1xf32> loc(#loc5259)
    %11480 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11481 = stablehlo.broadcast_in_dim %11480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11482 = stablehlo.subtract %11479, %11481 : tensor<1xf32> loc(#loc5260)
    %11483 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11484 = stablehlo.broadcast_in_dim %11483, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11485 = stablehlo.power %11484, %11482 : tensor<1xf32> loc(#loc5261)
    %11486 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11487 = stablehlo.broadcast_in_dim %11486, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11488 = stablehlo.subtract %11487, %11485 : tensor<1xf32> loc(#loc5260)
    %11489 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11490 = stablehlo.broadcast_in_dim %11489, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11491 = stablehlo.multiply %11490, %11488 : tensor<1xf32> loc(#loc5236)
    %11492 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11493 = stablehlo.broadcast_in_dim %11492, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11494 = stablehlo.power %11493, %11479 : tensor<1xf32> loc(#loc5261)
    %11495 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11496 = stablehlo.broadcast_in_dim %11495, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11497 = stablehlo.subtract %11496, %11494 : tensor<1xf32> loc(#loc5260)
    %11498 = stablehlo.divide %11491, %11497 : tensor<1xf32> loc(#loc5253)
    %11499 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11500 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11501 = stablehlo.broadcast_in_dim %11500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11502 = stablehlo.add %11499, %11501 : tensor<1xf32> loc(#loc5259)
    %11503 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11504 = stablehlo.broadcast_in_dim %11503, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11505 = stablehlo.subtract %11502, %11504 : tensor<1xf32> loc(#loc5260)
    %11506 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11507 = stablehlo.broadcast_in_dim %11506, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11508 = stablehlo.power %11507, %11505 : tensor<1xf32> loc(#loc5261)
    %11509 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11510 = stablehlo.broadcast_in_dim %11509, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11511 = stablehlo.subtract %11510, %11508 : tensor<1xf32> loc(#loc5260)
    %11512 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11513 = stablehlo.broadcast_in_dim %11512, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11514 = stablehlo.multiply %11513, %11511 : tensor<1xf32> loc(#loc5236)
    %11515 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11516 = stablehlo.broadcast_in_dim %11515, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11517 = stablehlo.power %11516, %11502 : tensor<1xf32> loc(#loc5261)
    %11518 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11519 = stablehlo.broadcast_in_dim %11518, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11520 = stablehlo.subtract %11519, %11517 : tensor<1xf32> loc(#loc5260)
    %11521 = stablehlo.divide %11514, %11520 : tensor<1xf32> loc(#loc5253)
    %11522 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11523 = stablehlo.broadcast_in_dim %11522, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11524 = stablehlo.subtract %11523, %11498 : tensor<1xf32> loc(#loc5260)
    %11525 = stablehlo.broadcast_in_dim %11524, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11526 = stablehlo.broadcast_in_dim %11525, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %11527 = stablehlo.multiply %11526, %4664 : tensor<1x64xf32> loc(#loc5262)
    %11528 = stablehlo.broadcast_in_dim %11498, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11529 = stablehlo.broadcast_in_dim %11528, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %11530 = stablehlo.multiply %11529, %arg296 : tensor<1x64xf32> loc(#loc5262)
    %11531 = stablehlo.add %11527, %11530 : tensor<1x64xf32> loc(#loc5263)
    %11532 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11533 = stablehlo.broadcast_in_dim %11532, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11534 = stablehlo.subtract %11533, %11521 : tensor<1xf32> loc(#loc5260)
    %11535 = stablehlo.multiply %4664, %4664 : tensor<1x64xf32> loc(#loc5264)
    %11536 = stablehlo.broadcast_in_dim %11534, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11537 = stablehlo.broadcast_in_dim %11536, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %11538 = stablehlo.multiply %11537, %11535 : tensor<1x64xf32> loc(#loc5262)
    %11539 = stablehlo.broadcast_in_dim %11521, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11540 = stablehlo.broadcast_in_dim %11539, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %11541 = stablehlo.multiply %11540, %arg472 : tensor<1x64xf32> loc(#loc5262)
    %11542 = stablehlo.add %11538, %11541 : tensor<1x64xf32> loc(#loc5263)
    %11543 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11544 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11545 = stablehlo.broadcast_in_dim %11544, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11546 = stablehlo.add %11543, %11545 : tensor<1xf32> loc(#loc5259)
    %11547 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11548 = stablehlo.broadcast_in_dim %11547, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11549 = stablehlo.subtract %11546, %11548 : tensor<1xf32> loc(#loc5260)
    %11550 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11551 = stablehlo.broadcast_in_dim %11550, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11552 = stablehlo.power %11551, %11549 : tensor<1xf32> loc(#loc5261)
    %11553 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11554 = stablehlo.broadcast_in_dim %11553, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11555 = stablehlo.subtract %11554, %11552 : tensor<1xf32> loc(#loc5260)
    %11556 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11557 = stablehlo.broadcast_in_dim %11556, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11558 = stablehlo.multiply %11557, %11555 : tensor<1xf32> loc(#loc5236)
    %11559 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11560 = stablehlo.broadcast_in_dim %11559, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11561 = stablehlo.power %11560, %11546 : tensor<1xf32> loc(#loc5261)
    %11562 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11563 = stablehlo.broadcast_in_dim %11562, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11564 = stablehlo.subtract %11563, %11561 : tensor<1xf32> loc(#loc5260)
    %11565 = stablehlo.divide %11558, %11564 : tensor<1xf32> loc(#loc5253)
    %11566 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11567 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11568 = stablehlo.broadcast_in_dim %11567, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11569 = stablehlo.add %11566, %11568 : tensor<1xf32> loc(#loc5259)
    %11570 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11571 = stablehlo.broadcast_in_dim %11570, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11572 = stablehlo.subtract %11569, %11571 : tensor<1xf32> loc(#loc5260)
    %11573 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11574 = stablehlo.broadcast_in_dim %11573, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11575 = stablehlo.power %11574, %11572 : tensor<1xf32> loc(#loc5261)
    %11576 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11577 = stablehlo.broadcast_in_dim %11576, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11578 = stablehlo.subtract %11577, %11575 : tensor<1xf32> loc(#loc5260)
    %11579 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11580 = stablehlo.broadcast_in_dim %11579, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11581 = stablehlo.multiply %11580, %11578 : tensor<1xf32> loc(#loc5236)
    %11582 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11583 = stablehlo.broadcast_in_dim %11582, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11584 = stablehlo.power %11583, %11569 : tensor<1xf32> loc(#loc5261)
    %11585 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11586 = stablehlo.broadcast_in_dim %11585, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11587 = stablehlo.subtract %11586, %11584 : tensor<1xf32> loc(#loc5260)
    %11588 = stablehlo.divide %11581, %11587 : tensor<1xf32> loc(#loc5253)
    %11589 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11590 = stablehlo.broadcast_in_dim %11589, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11591 = stablehlo.subtract %11590, %11565 : tensor<1xf32> loc(#loc5260)
    %11592 = stablehlo.broadcast_in_dim %11591, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %11593 = stablehlo.broadcast_in_dim %11592, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %11594 = stablehlo.multiply %11593, %4668 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %11595 = stablehlo.broadcast_in_dim %11565, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %11596 = stablehlo.broadcast_in_dim %11595, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %11597 = stablehlo.multiply %11596, %arg297 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %11598 = stablehlo.add %11594, %11597 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %11599 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11600 = stablehlo.broadcast_in_dim %11599, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11601 = stablehlo.subtract %11600, %11588 : tensor<1xf32> loc(#loc5260)
    %11602 = stablehlo.multiply %4668, %4668 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %11603 = stablehlo.broadcast_in_dim %11601, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %11604 = stablehlo.broadcast_in_dim %11603, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %11605 = stablehlo.multiply %11604, %11602 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %11606 = stablehlo.broadcast_in_dim %11588, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %11607 = stablehlo.broadcast_in_dim %11606, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %11608 = stablehlo.multiply %11607, %arg473 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %11609 = stablehlo.add %11605, %11608 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %11610 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11611 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11612 = stablehlo.broadcast_in_dim %11611, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11613 = stablehlo.add %11610, %11612 : tensor<1xf32> loc(#loc5259)
    %11614 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11615 = stablehlo.broadcast_in_dim %11614, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11616 = stablehlo.subtract %11613, %11615 : tensor<1xf32> loc(#loc5260)
    %11617 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11618 = stablehlo.broadcast_in_dim %11617, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11619 = stablehlo.power %11618, %11616 : tensor<1xf32> loc(#loc5261)
    %11620 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11621 = stablehlo.broadcast_in_dim %11620, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11622 = stablehlo.subtract %11621, %11619 : tensor<1xf32> loc(#loc5260)
    %11623 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11624 = stablehlo.broadcast_in_dim %11623, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11625 = stablehlo.multiply %11624, %11622 : tensor<1xf32> loc(#loc5236)
    %11626 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11627 = stablehlo.broadcast_in_dim %11626, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11628 = stablehlo.power %11627, %11613 : tensor<1xf32> loc(#loc5261)
    %11629 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11630 = stablehlo.broadcast_in_dim %11629, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11631 = stablehlo.subtract %11630, %11628 : tensor<1xf32> loc(#loc5260)
    %11632 = stablehlo.divide %11625, %11631 : tensor<1xf32> loc(#loc5253)
    %11633 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11634 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11635 = stablehlo.broadcast_in_dim %11634, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11636 = stablehlo.add %11633, %11635 : tensor<1xf32> loc(#loc5259)
    %11637 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11638 = stablehlo.broadcast_in_dim %11637, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11639 = stablehlo.subtract %11636, %11638 : tensor<1xf32> loc(#loc5260)
    %11640 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11641 = stablehlo.broadcast_in_dim %11640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11642 = stablehlo.power %11641, %11639 : tensor<1xf32> loc(#loc5261)
    %11643 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11644 = stablehlo.broadcast_in_dim %11643, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11645 = stablehlo.subtract %11644, %11642 : tensor<1xf32> loc(#loc5260)
    %11646 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11647 = stablehlo.broadcast_in_dim %11646, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11648 = stablehlo.multiply %11647, %11645 : tensor<1xf32> loc(#loc5236)
    %11649 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11650 = stablehlo.broadcast_in_dim %11649, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11651 = stablehlo.power %11650, %11636 : tensor<1xf32> loc(#loc5261)
    %11652 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11653 = stablehlo.broadcast_in_dim %11652, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11654 = stablehlo.subtract %11653, %11651 : tensor<1xf32> loc(#loc5260)
    %11655 = stablehlo.divide %11648, %11654 : tensor<1xf32> loc(#loc5253)
    %11656 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11657 = stablehlo.broadcast_in_dim %11656, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11658 = stablehlo.subtract %11657, %11632 : tensor<1xf32> loc(#loc5260)
    %11659 = stablehlo.broadcast_in_dim %11658, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11660 = stablehlo.broadcast_in_dim %11659, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %11661 = stablehlo.multiply %11660, %4672 : tensor<1x8192xf32> loc(#loc5262)
    %11662 = stablehlo.broadcast_in_dim %11632, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11663 = stablehlo.broadcast_in_dim %11662, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %11664 = stablehlo.multiply %11663, %arg298 : tensor<1x8192xf32> loc(#loc5262)
    %11665 = stablehlo.add %11661, %11664 : tensor<1x8192xf32> loc(#loc5263)
    %11666 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11667 = stablehlo.broadcast_in_dim %11666, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11668 = stablehlo.subtract %11667, %11655 : tensor<1xf32> loc(#loc5260)
    %11669 = stablehlo.multiply %4672, %4672 : tensor<1x8192xf32> loc(#loc5264)
    %11670 = stablehlo.broadcast_in_dim %11668, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11671 = stablehlo.broadcast_in_dim %11670, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %11672 = stablehlo.multiply %11671, %11669 : tensor<1x8192xf32> loc(#loc5262)
    %11673 = stablehlo.broadcast_in_dim %11655, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11674 = stablehlo.broadcast_in_dim %11673, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %11675 = stablehlo.multiply %11674, %arg474 : tensor<1x8192xf32> loc(#loc5262)
    %11676 = stablehlo.add %11672, %11675 : tensor<1x8192xf32> loc(#loc5263)
    %11677 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11678 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11679 = stablehlo.broadcast_in_dim %11678, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11680 = stablehlo.add %11677, %11679 : tensor<1xf32> loc(#loc5259)
    %11681 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11682 = stablehlo.broadcast_in_dim %11681, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11683 = stablehlo.subtract %11680, %11682 : tensor<1xf32> loc(#loc5260)
    %11684 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11685 = stablehlo.broadcast_in_dim %11684, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11686 = stablehlo.power %11685, %11683 : tensor<1xf32> loc(#loc5261)
    %11687 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11688 = stablehlo.broadcast_in_dim %11687, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11689 = stablehlo.subtract %11688, %11686 : tensor<1xf32> loc(#loc5260)
    %11690 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11691 = stablehlo.broadcast_in_dim %11690, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11692 = stablehlo.multiply %11691, %11689 : tensor<1xf32> loc(#loc5236)
    %11693 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11694 = stablehlo.broadcast_in_dim %11693, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11695 = stablehlo.power %11694, %11680 : tensor<1xf32> loc(#loc5261)
    %11696 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11697 = stablehlo.broadcast_in_dim %11696, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11698 = stablehlo.subtract %11697, %11695 : tensor<1xf32> loc(#loc5260)
    %11699 = stablehlo.divide %11692, %11698 : tensor<1xf32> loc(#loc5253)
    %11700 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11701 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11702 = stablehlo.broadcast_in_dim %11701, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11703 = stablehlo.add %11700, %11702 : tensor<1xf32> loc(#loc5259)
    %11704 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11705 = stablehlo.broadcast_in_dim %11704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11706 = stablehlo.subtract %11703, %11705 : tensor<1xf32> loc(#loc5260)
    %11707 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11708 = stablehlo.broadcast_in_dim %11707, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11709 = stablehlo.power %11708, %11706 : tensor<1xf32> loc(#loc5261)
    %11710 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11711 = stablehlo.broadcast_in_dim %11710, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11712 = stablehlo.subtract %11711, %11709 : tensor<1xf32> loc(#loc5260)
    %11713 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11714 = stablehlo.broadcast_in_dim %11713, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11715 = stablehlo.multiply %11714, %11712 : tensor<1xf32> loc(#loc5236)
    %11716 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11717 = stablehlo.broadcast_in_dim %11716, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11718 = stablehlo.power %11717, %11703 : tensor<1xf32> loc(#loc5261)
    %11719 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11720 = stablehlo.broadcast_in_dim %11719, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11721 = stablehlo.subtract %11720, %11718 : tensor<1xf32> loc(#loc5260)
    %11722 = stablehlo.divide %11715, %11721 : tensor<1xf32> loc(#loc5253)
    %11723 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11724 = stablehlo.broadcast_in_dim %11723, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11725 = stablehlo.subtract %11724, %11699 : tensor<1xf32> loc(#loc5260)
    %11726 = stablehlo.broadcast_in_dim %11725, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11727 = stablehlo.broadcast_in_dim %11726, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %11728 = stablehlo.multiply %11727, %4676 : tensor<1x2048x8192xf32> loc(#loc5262)
    %11729 = stablehlo.broadcast_in_dim %11699, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11730 = stablehlo.broadcast_in_dim %11729, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %11731 = stablehlo.multiply %11730, %arg299 : tensor<1x2048x8192xf32> loc(#loc5262)
    %11732 = stablehlo.add %11728, %11731 : tensor<1x2048x8192xf32> loc(#loc5263)
    %11733 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11734 = stablehlo.broadcast_in_dim %11733, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11735 = stablehlo.subtract %11734, %11722 : tensor<1xf32> loc(#loc5260)
    %11736 = stablehlo.multiply %4676, %4676 : tensor<1x2048x8192xf32> loc(#loc5264)
    %11737 = stablehlo.broadcast_in_dim %11735, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11738 = stablehlo.broadcast_in_dim %11737, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %11739 = stablehlo.multiply %11738, %11736 : tensor<1x2048x8192xf32> loc(#loc5262)
    %11740 = stablehlo.broadcast_in_dim %11722, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11741 = stablehlo.broadcast_in_dim %11740, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %11742 = stablehlo.multiply %11741, %arg475 : tensor<1x2048x8192xf32> loc(#loc5262)
    %11743 = stablehlo.add %11739, %11742 : tensor<1x2048x8192xf32> loc(#loc5263)
    %11744 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11745 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11746 = stablehlo.broadcast_in_dim %11745, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11747 = stablehlo.add %11744, %11746 : tensor<1xf32> loc(#loc5259)
    %11748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11749 = stablehlo.broadcast_in_dim %11748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11750 = stablehlo.subtract %11747, %11749 : tensor<1xf32> loc(#loc5260)
    %11751 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11752 = stablehlo.broadcast_in_dim %11751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11753 = stablehlo.power %11752, %11750 : tensor<1xf32> loc(#loc5261)
    %11754 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11755 = stablehlo.broadcast_in_dim %11754, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11756 = stablehlo.subtract %11755, %11753 : tensor<1xf32> loc(#loc5260)
    %11757 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11758 = stablehlo.broadcast_in_dim %11757, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11759 = stablehlo.multiply %11758, %11756 : tensor<1xf32> loc(#loc5236)
    %11760 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11761 = stablehlo.broadcast_in_dim %11760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11762 = stablehlo.power %11761, %11747 : tensor<1xf32> loc(#loc5261)
    %11763 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11764 = stablehlo.broadcast_in_dim %11763, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11765 = stablehlo.subtract %11764, %11762 : tensor<1xf32> loc(#loc5260)
    %11766 = stablehlo.divide %11759, %11765 : tensor<1xf32> loc(#loc5253)
    %11767 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11768 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11769 = stablehlo.broadcast_in_dim %11768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11770 = stablehlo.add %11767, %11769 : tensor<1xf32> loc(#loc5259)
    %11771 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11772 = stablehlo.broadcast_in_dim %11771, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11773 = stablehlo.subtract %11770, %11772 : tensor<1xf32> loc(#loc5260)
    %11774 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11775 = stablehlo.broadcast_in_dim %11774, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11776 = stablehlo.power %11775, %11773 : tensor<1xf32> loc(#loc5261)
    %11777 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11778 = stablehlo.broadcast_in_dim %11777, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11779 = stablehlo.subtract %11778, %11776 : tensor<1xf32> loc(#loc5260)
    %11780 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11781 = stablehlo.broadcast_in_dim %11780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11782 = stablehlo.multiply %11781, %11779 : tensor<1xf32> loc(#loc5236)
    %11783 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11784 = stablehlo.broadcast_in_dim %11783, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11785 = stablehlo.power %11784, %11770 : tensor<1xf32> loc(#loc5261)
    %11786 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11787 = stablehlo.broadcast_in_dim %11786, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11788 = stablehlo.subtract %11787, %11785 : tensor<1xf32> loc(#loc5260)
    %11789 = stablehlo.divide %11782, %11788 : tensor<1xf32> loc(#loc5253)
    %11790 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11791 = stablehlo.broadcast_in_dim %11790, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11792 = stablehlo.subtract %11791, %11766 : tensor<1xf32> loc(#loc5260)
    %11793 = stablehlo.broadcast_in_dim %11792, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11794 = stablehlo.broadcast_in_dim %11793, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11795 = stablehlo.multiply %11794, %4680 : tensor<1x2048xf32> loc(#loc5262)
    %11796 = stablehlo.broadcast_in_dim %11766, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11797 = stablehlo.broadcast_in_dim %11796, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11798 = stablehlo.multiply %11797, %arg300 : tensor<1x2048xf32> loc(#loc5262)
    %11799 = stablehlo.add %11795, %11798 : tensor<1x2048xf32> loc(#loc5263)
    %11800 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11801 = stablehlo.broadcast_in_dim %11800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11802 = stablehlo.subtract %11801, %11789 : tensor<1xf32> loc(#loc5260)
    %11803 = stablehlo.multiply %4680, %4680 : tensor<1x2048xf32> loc(#loc5264)
    %11804 = stablehlo.broadcast_in_dim %11802, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11805 = stablehlo.broadcast_in_dim %11804, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11806 = stablehlo.multiply %11805, %11803 : tensor<1x2048xf32> loc(#loc5262)
    %11807 = stablehlo.broadcast_in_dim %11789, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11808 = stablehlo.broadcast_in_dim %11807, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11809 = stablehlo.multiply %11808, %arg476 : tensor<1x2048xf32> loc(#loc5262)
    %11810 = stablehlo.add %11806, %11809 : tensor<1x2048xf32> loc(#loc5263)
    %11811 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11812 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11813 = stablehlo.broadcast_in_dim %11812, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11814 = stablehlo.add %11811, %11813 : tensor<1xf32> loc(#loc5259)
    %11815 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11816 = stablehlo.broadcast_in_dim %11815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11817 = stablehlo.subtract %11814, %11816 : tensor<1xf32> loc(#loc5260)
    %11818 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11819 = stablehlo.broadcast_in_dim %11818, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11820 = stablehlo.power %11819, %11817 : tensor<1xf32> loc(#loc5261)
    %11821 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11822 = stablehlo.broadcast_in_dim %11821, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11823 = stablehlo.subtract %11822, %11820 : tensor<1xf32> loc(#loc5260)
    %11824 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11825 = stablehlo.broadcast_in_dim %11824, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11826 = stablehlo.multiply %11825, %11823 : tensor<1xf32> loc(#loc5236)
    %11827 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11828 = stablehlo.broadcast_in_dim %11827, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11829 = stablehlo.power %11828, %11814 : tensor<1xf32> loc(#loc5261)
    %11830 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11831 = stablehlo.broadcast_in_dim %11830, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11832 = stablehlo.subtract %11831, %11829 : tensor<1xf32> loc(#loc5260)
    %11833 = stablehlo.divide %11826, %11832 : tensor<1xf32> loc(#loc5253)
    %11834 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11835 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11836 = stablehlo.broadcast_in_dim %11835, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11837 = stablehlo.add %11834, %11836 : tensor<1xf32> loc(#loc5259)
    %11838 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11839 = stablehlo.broadcast_in_dim %11838, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11840 = stablehlo.subtract %11837, %11839 : tensor<1xf32> loc(#loc5260)
    %11841 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11842 = stablehlo.broadcast_in_dim %11841, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11843 = stablehlo.power %11842, %11840 : tensor<1xf32> loc(#loc5261)
    %11844 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11845 = stablehlo.broadcast_in_dim %11844, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11846 = stablehlo.subtract %11845, %11843 : tensor<1xf32> loc(#loc5260)
    %11847 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11848 = stablehlo.broadcast_in_dim %11847, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11849 = stablehlo.multiply %11848, %11846 : tensor<1xf32> loc(#loc5236)
    %11850 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11851 = stablehlo.broadcast_in_dim %11850, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11852 = stablehlo.power %11851, %11837 : tensor<1xf32> loc(#loc5261)
    %11853 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11854 = stablehlo.broadcast_in_dim %11853, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11855 = stablehlo.subtract %11854, %11852 : tensor<1xf32> loc(#loc5260)
    %11856 = stablehlo.divide %11849, %11855 : tensor<1xf32> loc(#loc5253)
    %11857 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11858 = stablehlo.broadcast_in_dim %11857, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11859 = stablehlo.subtract %11858, %11833 : tensor<1xf32> loc(#loc5260)
    %11860 = stablehlo.broadcast_in_dim %11859, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11861 = stablehlo.broadcast_in_dim %11860, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11862 = stablehlo.multiply %11861, %4684 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11863 = stablehlo.broadcast_in_dim %11833, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11864 = stablehlo.broadcast_in_dim %11863, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11865 = stablehlo.multiply %11864, %arg301 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11866 = stablehlo.add %11862, %11865 : tensor<1x8192x2048xf32> loc(#loc5263)
    %11867 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11868 = stablehlo.broadcast_in_dim %11867, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11869 = stablehlo.subtract %11868, %11856 : tensor<1xf32> loc(#loc5260)
    %11870 = stablehlo.multiply %4684, %4684 : tensor<1x8192x2048xf32> loc(#loc5264)
    %11871 = stablehlo.broadcast_in_dim %11869, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11872 = stablehlo.broadcast_in_dim %11871, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11873 = stablehlo.multiply %11872, %11870 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11874 = stablehlo.broadcast_in_dim %11856, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %11875 = stablehlo.broadcast_in_dim %11874, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %11876 = stablehlo.multiply %11875, %arg477 : tensor<1x8192x2048xf32> loc(#loc5262)
    %11877 = stablehlo.add %11873, %11876 : tensor<1x8192x2048xf32> loc(#loc5263)
    %11878 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11879 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11880 = stablehlo.broadcast_in_dim %11879, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11881 = stablehlo.add %11878, %11880 : tensor<1xf32> loc(#loc5259)
    %11882 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11883 = stablehlo.broadcast_in_dim %11882, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11884 = stablehlo.subtract %11881, %11883 : tensor<1xf32> loc(#loc5260)
    %11885 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11886 = stablehlo.broadcast_in_dim %11885, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11887 = stablehlo.power %11886, %11884 : tensor<1xf32> loc(#loc5261)
    %11888 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11889 = stablehlo.broadcast_in_dim %11888, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11890 = stablehlo.subtract %11889, %11887 : tensor<1xf32> loc(#loc5260)
    %11891 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11892 = stablehlo.broadcast_in_dim %11891, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11893 = stablehlo.multiply %11892, %11890 : tensor<1xf32> loc(#loc5236)
    %11894 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11895 = stablehlo.broadcast_in_dim %11894, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11896 = stablehlo.power %11895, %11881 : tensor<1xf32> loc(#loc5261)
    %11897 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11898 = stablehlo.broadcast_in_dim %11897, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11899 = stablehlo.subtract %11898, %11896 : tensor<1xf32> loc(#loc5260)
    %11900 = stablehlo.divide %11893, %11899 : tensor<1xf32> loc(#loc5253)
    %11901 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11902 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11903 = stablehlo.broadcast_in_dim %11902, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11904 = stablehlo.add %11901, %11903 : tensor<1xf32> loc(#loc5259)
    %11905 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11906 = stablehlo.broadcast_in_dim %11905, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11907 = stablehlo.subtract %11904, %11906 : tensor<1xf32> loc(#loc5260)
    %11908 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11909 = stablehlo.broadcast_in_dim %11908, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11910 = stablehlo.power %11909, %11907 : tensor<1xf32> loc(#loc5261)
    %11911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11912 = stablehlo.broadcast_in_dim %11911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11913 = stablehlo.subtract %11912, %11910 : tensor<1xf32> loc(#loc5260)
    %11914 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11915 = stablehlo.broadcast_in_dim %11914, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11916 = stablehlo.multiply %11915, %11913 : tensor<1xf32> loc(#loc5236)
    %11917 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11918 = stablehlo.broadcast_in_dim %11917, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11919 = stablehlo.power %11918, %11904 : tensor<1xf32> loc(#loc5261)
    %11920 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11921 = stablehlo.broadcast_in_dim %11920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11922 = stablehlo.subtract %11921, %11919 : tensor<1xf32> loc(#loc5260)
    %11923 = stablehlo.divide %11916, %11922 : tensor<1xf32> loc(#loc5253)
    %11924 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11925 = stablehlo.broadcast_in_dim %11924, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11926 = stablehlo.subtract %11925, %11900 : tensor<1xf32> loc(#loc5260)
    %11927 = stablehlo.broadcast_in_dim %11926, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11928 = stablehlo.broadcast_in_dim %11927, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11929 = stablehlo.multiply %11928, %4688 : tensor<1x2048xf32> loc(#loc5262)
    %11930 = stablehlo.broadcast_in_dim %11900, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11931 = stablehlo.broadcast_in_dim %11930, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11932 = stablehlo.multiply %11931, %arg302 : tensor<1x2048xf32> loc(#loc5262)
    %11933 = stablehlo.add %11929, %11932 : tensor<1x2048xf32> loc(#loc5263)
    %11934 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11935 = stablehlo.broadcast_in_dim %11934, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11936 = stablehlo.subtract %11935, %11923 : tensor<1xf32> loc(#loc5260)
    %11937 = stablehlo.multiply %4688, %4688 : tensor<1x2048xf32> loc(#loc5264)
    %11938 = stablehlo.broadcast_in_dim %11936, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11939 = stablehlo.broadcast_in_dim %11938, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11940 = stablehlo.multiply %11939, %11937 : tensor<1x2048xf32> loc(#loc5262)
    %11941 = stablehlo.broadcast_in_dim %11923, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11942 = stablehlo.broadcast_in_dim %11941, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11943 = stablehlo.multiply %11942, %arg478 : tensor<1x2048xf32> loc(#loc5262)
    %11944 = stablehlo.add %11940, %11943 : tensor<1x2048xf32> loc(#loc5263)
    %11945 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11946 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11947 = stablehlo.broadcast_in_dim %11946, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11948 = stablehlo.add %11945, %11947 : tensor<1xf32> loc(#loc5259)
    %11949 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11950 = stablehlo.broadcast_in_dim %11949, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11951 = stablehlo.subtract %11948, %11950 : tensor<1xf32> loc(#loc5260)
    %11952 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11953 = stablehlo.broadcast_in_dim %11952, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11954 = stablehlo.power %11953, %11951 : tensor<1xf32> loc(#loc5261)
    %11955 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11956 = stablehlo.broadcast_in_dim %11955, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11957 = stablehlo.subtract %11956, %11954 : tensor<1xf32> loc(#loc5260)
    %11958 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11959 = stablehlo.broadcast_in_dim %11958, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11960 = stablehlo.multiply %11959, %11957 : tensor<1xf32> loc(#loc5236)
    %11961 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %11962 = stablehlo.broadcast_in_dim %11961, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11963 = stablehlo.power %11962, %11948 : tensor<1xf32> loc(#loc5261)
    %11964 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11965 = stablehlo.broadcast_in_dim %11964, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11966 = stablehlo.subtract %11965, %11963 : tensor<1xf32> loc(#loc5260)
    %11967 = stablehlo.divide %11960, %11966 : tensor<1xf32> loc(#loc5253)
    %11968 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %11969 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11970 = stablehlo.broadcast_in_dim %11969, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %11971 = stablehlo.add %11968, %11970 : tensor<1xf32> loc(#loc5259)
    %11972 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11973 = stablehlo.broadcast_in_dim %11972, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11974 = stablehlo.subtract %11971, %11973 : tensor<1xf32> loc(#loc5260)
    %11975 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11976 = stablehlo.broadcast_in_dim %11975, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11977 = stablehlo.power %11976, %11974 : tensor<1xf32> loc(#loc5261)
    %11978 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11979 = stablehlo.broadcast_in_dim %11978, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11980 = stablehlo.subtract %11979, %11977 : tensor<1xf32> loc(#loc5260)
    %11981 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11982 = stablehlo.broadcast_in_dim %11981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %11983 = stablehlo.multiply %11982, %11980 : tensor<1xf32> loc(#loc5236)
    %11984 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %11985 = stablehlo.broadcast_in_dim %11984, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %11986 = stablehlo.power %11985, %11971 : tensor<1xf32> loc(#loc5261)
    %11987 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11988 = stablehlo.broadcast_in_dim %11987, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11989 = stablehlo.subtract %11988, %11986 : tensor<1xf32> loc(#loc5260)
    %11990 = stablehlo.divide %11983, %11989 : tensor<1xf32> loc(#loc5253)
    %11991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %11992 = stablehlo.broadcast_in_dim %11991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %11993 = stablehlo.subtract %11992, %11967 : tensor<1xf32> loc(#loc5260)
    %11994 = stablehlo.broadcast_in_dim %11993, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11995 = stablehlo.broadcast_in_dim %11994, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11996 = stablehlo.multiply %11995, %4692 : tensor<1x2048xf32> loc(#loc5262)
    %11997 = stablehlo.broadcast_in_dim %11967, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %11998 = stablehlo.broadcast_in_dim %11997, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %11999 = stablehlo.multiply %11998, %arg303 : tensor<1x2048xf32> loc(#loc5262)
    %12000 = stablehlo.add %11996, %11999 : tensor<1x2048xf32> loc(#loc5263)
    %12001 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12002 = stablehlo.broadcast_in_dim %12001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12003 = stablehlo.subtract %12002, %11990 : tensor<1xf32> loc(#loc5260)
    %12004 = stablehlo.multiply %4692, %4692 : tensor<1x2048xf32> loc(#loc5264)
    %12005 = stablehlo.broadcast_in_dim %12003, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12006 = stablehlo.broadcast_in_dim %12005, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12007 = stablehlo.multiply %12006, %12004 : tensor<1x2048xf32> loc(#loc5262)
    %12008 = stablehlo.broadcast_in_dim %11990, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12009 = stablehlo.broadcast_in_dim %12008, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12010 = stablehlo.multiply %12009, %arg479 : tensor<1x2048xf32> loc(#loc5262)
    %12011 = stablehlo.add %12007, %12010 : tensor<1x2048xf32> loc(#loc5263)
    %12012 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12013 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12014 = stablehlo.broadcast_in_dim %12013, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12015 = stablehlo.add %12012, %12014 : tensor<1xf32> loc(#loc5259)
    %12016 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12017 = stablehlo.broadcast_in_dim %12016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12018 = stablehlo.subtract %12015, %12017 : tensor<1xf32> loc(#loc5260)
    %12019 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12020 = stablehlo.broadcast_in_dim %12019, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12021 = stablehlo.power %12020, %12018 : tensor<1xf32> loc(#loc5261)
    %12022 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12023 = stablehlo.broadcast_in_dim %12022, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12024 = stablehlo.subtract %12023, %12021 : tensor<1xf32> loc(#loc5260)
    %12025 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12026 = stablehlo.broadcast_in_dim %12025, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12027 = stablehlo.multiply %12026, %12024 : tensor<1xf32> loc(#loc5236)
    %12028 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12029 = stablehlo.broadcast_in_dim %12028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12030 = stablehlo.power %12029, %12015 : tensor<1xf32> loc(#loc5261)
    %12031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12032 = stablehlo.broadcast_in_dim %12031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12033 = stablehlo.subtract %12032, %12030 : tensor<1xf32> loc(#loc5260)
    %12034 = stablehlo.divide %12027, %12033 : tensor<1xf32> loc(#loc5253)
    %12035 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12036 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12037 = stablehlo.broadcast_in_dim %12036, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12038 = stablehlo.add %12035, %12037 : tensor<1xf32> loc(#loc5259)
    %12039 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12040 = stablehlo.broadcast_in_dim %12039, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12041 = stablehlo.subtract %12038, %12040 : tensor<1xf32> loc(#loc5260)
    %12042 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12043 = stablehlo.broadcast_in_dim %12042, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12044 = stablehlo.power %12043, %12041 : tensor<1xf32> loc(#loc5261)
    %12045 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12046 = stablehlo.broadcast_in_dim %12045, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12047 = stablehlo.subtract %12046, %12044 : tensor<1xf32> loc(#loc5260)
    %12048 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12049 = stablehlo.broadcast_in_dim %12048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12050 = stablehlo.multiply %12049, %12047 : tensor<1xf32> loc(#loc5236)
    %12051 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12052 = stablehlo.broadcast_in_dim %12051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12053 = stablehlo.power %12052, %12038 : tensor<1xf32> loc(#loc5261)
    %12054 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12055 = stablehlo.broadcast_in_dim %12054, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12056 = stablehlo.subtract %12055, %12053 : tensor<1xf32> loc(#loc5260)
    %12057 = stablehlo.divide %12050, %12056 : tensor<1xf32> loc(#loc5253)
    %12058 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12059 = stablehlo.broadcast_in_dim %12058, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12060 = stablehlo.subtract %12059, %12034 : tensor<1xf32> loc(#loc5260)
    %12061 = stablehlo.broadcast_in_dim %12060, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12062 = stablehlo.broadcast_in_dim %12061, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12063 = stablehlo.multiply %12062, %4696 : tensor<1x2048xf32> loc(#loc5262)
    %12064 = stablehlo.broadcast_in_dim %12034, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12065 = stablehlo.broadcast_in_dim %12064, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12066 = stablehlo.multiply %12065, %arg304 : tensor<1x2048xf32> loc(#loc5262)
    %12067 = stablehlo.add %12063, %12066 : tensor<1x2048xf32> loc(#loc5263)
    %12068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12069 = stablehlo.broadcast_in_dim %12068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12070 = stablehlo.subtract %12069, %12057 : tensor<1xf32> loc(#loc5260)
    %12071 = stablehlo.multiply %4696, %4696 : tensor<1x2048xf32> loc(#loc5264)
    %12072 = stablehlo.broadcast_in_dim %12070, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12073 = stablehlo.broadcast_in_dim %12072, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12074 = stablehlo.multiply %12073, %12071 : tensor<1x2048xf32> loc(#loc5262)
    %12075 = stablehlo.broadcast_in_dim %12057, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12076 = stablehlo.broadcast_in_dim %12075, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12077 = stablehlo.multiply %12076, %arg480 : tensor<1x2048xf32> loc(#loc5262)
    %12078 = stablehlo.add %12074, %12077 : tensor<1x2048xf32> loc(#loc5263)
    %12079 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12080 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12081 = stablehlo.broadcast_in_dim %12080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12082 = stablehlo.add %12079, %12081 : tensor<1xf32> loc(#loc5259)
    %12083 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12084 = stablehlo.broadcast_in_dim %12083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12085 = stablehlo.subtract %12082, %12084 : tensor<1xf32> loc(#loc5260)
    %12086 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12087 = stablehlo.broadcast_in_dim %12086, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12088 = stablehlo.power %12087, %12085 : tensor<1xf32> loc(#loc5261)
    %12089 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12090 = stablehlo.broadcast_in_dim %12089, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12091 = stablehlo.subtract %12090, %12088 : tensor<1xf32> loc(#loc5260)
    %12092 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12093 = stablehlo.broadcast_in_dim %12092, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12094 = stablehlo.multiply %12093, %12091 : tensor<1xf32> loc(#loc5236)
    %12095 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12096 = stablehlo.broadcast_in_dim %12095, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12097 = stablehlo.power %12096, %12082 : tensor<1xf32> loc(#loc5261)
    %12098 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12099 = stablehlo.broadcast_in_dim %12098, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12100 = stablehlo.subtract %12099, %12097 : tensor<1xf32> loc(#loc5260)
    %12101 = stablehlo.divide %12094, %12100 : tensor<1xf32> loc(#loc5253)
    %12102 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12103 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12104 = stablehlo.broadcast_in_dim %12103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12105 = stablehlo.add %12102, %12104 : tensor<1xf32> loc(#loc5259)
    %12106 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12107 = stablehlo.broadcast_in_dim %12106, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12108 = stablehlo.subtract %12105, %12107 : tensor<1xf32> loc(#loc5260)
    %12109 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12110 = stablehlo.broadcast_in_dim %12109, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12111 = stablehlo.power %12110, %12108 : tensor<1xf32> loc(#loc5261)
    %12112 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12113 = stablehlo.broadcast_in_dim %12112, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12114 = stablehlo.subtract %12113, %12111 : tensor<1xf32> loc(#loc5260)
    %12115 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12116 = stablehlo.broadcast_in_dim %12115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12117 = stablehlo.multiply %12116, %12114 : tensor<1xf32> loc(#loc5236)
    %12118 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12119 = stablehlo.broadcast_in_dim %12118, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12120 = stablehlo.power %12119, %12105 : tensor<1xf32> loc(#loc5261)
    %12121 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12122 = stablehlo.broadcast_in_dim %12121, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12123 = stablehlo.subtract %12122, %12120 : tensor<1xf32> loc(#loc5260)
    %12124 = stablehlo.divide %12117, %12123 : tensor<1xf32> loc(#loc5253)
    %12125 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12126 = stablehlo.broadcast_in_dim %12125, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12127 = stablehlo.subtract %12126, %12101 : tensor<1xf32> loc(#loc5260)
    %12128 = stablehlo.broadcast_in_dim %12127, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12129 = stablehlo.broadcast_in_dim %12128, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12130 = stablehlo.multiply %12129, %4700 : tensor<1x2048xf32> loc(#loc5262)
    %12131 = stablehlo.broadcast_in_dim %12101, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12132 = stablehlo.broadcast_in_dim %12131, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12133 = stablehlo.multiply %12132, %arg305 : tensor<1x2048xf32> loc(#loc5262)
    %12134 = stablehlo.add %12130, %12133 : tensor<1x2048xf32> loc(#loc5263)
    %12135 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12136 = stablehlo.broadcast_in_dim %12135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12137 = stablehlo.subtract %12136, %12124 : tensor<1xf32> loc(#loc5260)
    %12138 = stablehlo.multiply %4700, %4700 : tensor<1x2048xf32> loc(#loc5264)
    %12139 = stablehlo.broadcast_in_dim %12137, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12140 = stablehlo.broadcast_in_dim %12139, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12141 = stablehlo.multiply %12140, %12138 : tensor<1x2048xf32> loc(#loc5262)
    %12142 = stablehlo.broadcast_in_dim %12124, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12143 = stablehlo.broadcast_in_dim %12142, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12144 = stablehlo.multiply %12143, %arg481 : tensor<1x2048xf32> loc(#loc5262)
    %12145 = stablehlo.add %12141, %12144 : tensor<1x2048xf32> loc(#loc5263)
    %12146 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12147 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12148 = stablehlo.broadcast_in_dim %12147, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12149 = stablehlo.add %12146, %12148 : tensor<1xf32> loc(#loc5259)
    %12150 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12151 = stablehlo.broadcast_in_dim %12150, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12152 = stablehlo.subtract %12149, %12151 : tensor<1xf32> loc(#loc5260)
    %12153 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12154 = stablehlo.broadcast_in_dim %12153, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12155 = stablehlo.power %12154, %12152 : tensor<1xf32> loc(#loc5261)
    %12156 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12157 = stablehlo.broadcast_in_dim %12156, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12158 = stablehlo.subtract %12157, %12155 : tensor<1xf32> loc(#loc5260)
    %12159 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12160 = stablehlo.broadcast_in_dim %12159, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12161 = stablehlo.multiply %12160, %12158 : tensor<1xf32> loc(#loc5236)
    %12162 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12163 = stablehlo.broadcast_in_dim %12162, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12164 = stablehlo.power %12163, %12149 : tensor<1xf32> loc(#loc5261)
    %12165 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12166 = stablehlo.broadcast_in_dim %12165, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12167 = stablehlo.subtract %12166, %12164 : tensor<1xf32> loc(#loc5260)
    %12168 = stablehlo.divide %12161, %12167 : tensor<1xf32> loc(#loc5253)
    %12169 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12170 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12171 = stablehlo.broadcast_in_dim %12170, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12172 = stablehlo.add %12169, %12171 : tensor<1xf32> loc(#loc5259)
    %12173 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12174 = stablehlo.broadcast_in_dim %12173, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12175 = stablehlo.subtract %12172, %12174 : tensor<1xf32> loc(#loc5260)
    %12176 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12177 = stablehlo.broadcast_in_dim %12176, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12178 = stablehlo.power %12177, %12175 : tensor<1xf32> loc(#loc5261)
    %12179 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12180 = stablehlo.broadcast_in_dim %12179, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12181 = stablehlo.subtract %12180, %12178 : tensor<1xf32> loc(#loc5260)
    %12182 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12183 = stablehlo.broadcast_in_dim %12182, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12184 = stablehlo.multiply %12183, %12181 : tensor<1xf32> loc(#loc5236)
    %12185 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12186 = stablehlo.broadcast_in_dim %12185, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12187 = stablehlo.power %12186, %12172 : tensor<1xf32> loc(#loc5261)
    %12188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12189 = stablehlo.broadcast_in_dim %12188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12190 = stablehlo.subtract %12189, %12187 : tensor<1xf32> loc(#loc5260)
    %12191 = stablehlo.divide %12184, %12190 : tensor<1xf32> loc(#loc5253)
    %12192 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12193 = stablehlo.broadcast_in_dim %12192, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12194 = stablehlo.subtract %12193, %12168 : tensor<1xf32> loc(#loc5260)
    %12195 = stablehlo.broadcast_in_dim %12194, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12196 = stablehlo.broadcast_in_dim %12195, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12197 = stablehlo.multiply %12196, %4704 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12198 = stablehlo.broadcast_in_dim %12168, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12199 = stablehlo.broadcast_in_dim %12198, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12200 = stablehlo.multiply %12199, %arg306 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12201 = stablehlo.add %12197, %12200 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %12202 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12203 = stablehlo.broadcast_in_dim %12202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12204 = stablehlo.subtract %12203, %12191 : tensor<1xf32> loc(#loc5260)
    %12205 = stablehlo.multiply %4704, %4704 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %12206 = stablehlo.broadcast_in_dim %12204, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12207 = stablehlo.broadcast_in_dim %12206, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12208 = stablehlo.multiply %12207, %12205 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12209 = stablehlo.broadcast_in_dim %12191, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12210 = stablehlo.broadcast_in_dim %12209, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12211 = stablehlo.multiply %12210, %arg482 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12212 = stablehlo.add %12208, %12211 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %12213 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12214 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12215 = stablehlo.broadcast_in_dim %12214, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12216 = stablehlo.add %12213, %12215 : tensor<1xf32> loc(#loc5259)
    %12217 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12218 = stablehlo.broadcast_in_dim %12217, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12219 = stablehlo.subtract %12216, %12218 : tensor<1xf32> loc(#loc5260)
    %12220 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12221 = stablehlo.broadcast_in_dim %12220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12222 = stablehlo.power %12221, %12219 : tensor<1xf32> loc(#loc5261)
    %12223 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12224 = stablehlo.broadcast_in_dim %12223, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12225 = stablehlo.subtract %12224, %12222 : tensor<1xf32> loc(#loc5260)
    %12226 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12227 = stablehlo.broadcast_in_dim %12226, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12228 = stablehlo.multiply %12227, %12225 : tensor<1xf32> loc(#loc5236)
    %12229 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12230 = stablehlo.broadcast_in_dim %12229, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12231 = stablehlo.power %12230, %12216 : tensor<1xf32> loc(#loc5261)
    %12232 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12233 = stablehlo.broadcast_in_dim %12232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12234 = stablehlo.subtract %12233, %12231 : tensor<1xf32> loc(#loc5260)
    %12235 = stablehlo.divide %12228, %12234 : tensor<1xf32> loc(#loc5253)
    %12236 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12237 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12238 = stablehlo.broadcast_in_dim %12237, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12239 = stablehlo.add %12236, %12238 : tensor<1xf32> loc(#loc5259)
    %12240 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12241 = stablehlo.broadcast_in_dim %12240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12242 = stablehlo.subtract %12239, %12241 : tensor<1xf32> loc(#loc5260)
    %12243 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12244 = stablehlo.broadcast_in_dim %12243, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12245 = stablehlo.power %12244, %12242 : tensor<1xf32> loc(#loc5261)
    %12246 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12247 = stablehlo.broadcast_in_dim %12246, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12248 = stablehlo.subtract %12247, %12245 : tensor<1xf32> loc(#loc5260)
    %12249 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12250 = stablehlo.broadcast_in_dim %12249, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12251 = stablehlo.multiply %12250, %12248 : tensor<1xf32> loc(#loc5236)
    %12252 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12253 = stablehlo.broadcast_in_dim %12252, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12254 = stablehlo.power %12253, %12239 : tensor<1xf32> loc(#loc5261)
    %12255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12256 = stablehlo.broadcast_in_dim %12255, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12257 = stablehlo.subtract %12256, %12254 : tensor<1xf32> loc(#loc5260)
    %12258 = stablehlo.divide %12251, %12257 : tensor<1xf32> loc(#loc5253)
    %12259 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12260 = stablehlo.broadcast_in_dim %12259, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12261 = stablehlo.subtract %12260, %12235 : tensor<1xf32> loc(#loc5260)
    %12262 = stablehlo.broadcast_in_dim %12261, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12263 = stablehlo.broadcast_in_dim %12262, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %12264 = stablehlo.multiply %12263, %4708 : tensor<1x64xf32> loc(#loc5262)
    %12265 = stablehlo.broadcast_in_dim %12235, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12266 = stablehlo.broadcast_in_dim %12265, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %12267 = stablehlo.multiply %12266, %arg307 : tensor<1x64xf32> loc(#loc5262)
    %12268 = stablehlo.add %12264, %12267 : tensor<1x64xf32> loc(#loc5263)
    %12269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12270 = stablehlo.broadcast_in_dim %12269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12271 = stablehlo.subtract %12270, %12258 : tensor<1xf32> loc(#loc5260)
    %12272 = stablehlo.multiply %4708, %4708 : tensor<1x64xf32> loc(#loc5264)
    %12273 = stablehlo.broadcast_in_dim %12271, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12274 = stablehlo.broadcast_in_dim %12273, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %12275 = stablehlo.multiply %12274, %12272 : tensor<1x64xf32> loc(#loc5262)
    %12276 = stablehlo.broadcast_in_dim %12258, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12277 = stablehlo.broadcast_in_dim %12276, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %12278 = stablehlo.multiply %12277, %arg483 : tensor<1x64xf32> loc(#loc5262)
    %12279 = stablehlo.add %12275, %12278 : tensor<1x64xf32> loc(#loc5263)
    %12280 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12281 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12282 = stablehlo.broadcast_in_dim %12281, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12283 = stablehlo.add %12280, %12282 : tensor<1xf32> loc(#loc5259)
    %12284 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12285 = stablehlo.broadcast_in_dim %12284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12286 = stablehlo.subtract %12283, %12285 : tensor<1xf32> loc(#loc5260)
    %12287 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12288 = stablehlo.broadcast_in_dim %12287, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12289 = stablehlo.power %12288, %12286 : tensor<1xf32> loc(#loc5261)
    %12290 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12291 = stablehlo.broadcast_in_dim %12290, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12292 = stablehlo.subtract %12291, %12289 : tensor<1xf32> loc(#loc5260)
    %12293 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12294 = stablehlo.broadcast_in_dim %12293, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12295 = stablehlo.multiply %12294, %12292 : tensor<1xf32> loc(#loc5236)
    %12296 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12297 = stablehlo.broadcast_in_dim %12296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12298 = stablehlo.power %12297, %12283 : tensor<1xf32> loc(#loc5261)
    %12299 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12300 = stablehlo.broadcast_in_dim %12299, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12301 = stablehlo.subtract %12300, %12298 : tensor<1xf32> loc(#loc5260)
    %12302 = stablehlo.divide %12295, %12301 : tensor<1xf32> loc(#loc5253)
    %12303 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12304 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12305 = stablehlo.broadcast_in_dim %12304, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12306 = stablehlo.add %12303, %12305 : tensor<1xf32> loc(#loc5259)
    %12307 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12308 = stablehlo.broadcast_in_dim %12307, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12309 = stablehlo.subtract %12306, %12308 : tensor<1xf32> loc(#loc5260)
    %12310 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12311 = stablehlo.broadcast_in_dim %12310, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12312 = stablehlo.power %12311, %12309 : tensor<1xf32> loc(#loc5261)
    %12313 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12314 = stablehlo.broadcast_in_dim %12313, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12315 = stablehlo.subtract %12314, %12312 : tensor<1xf32> loc(#loc5260)
    %12316 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12317 = stablehlo.broadcast_in_dim %12316, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12318 = stablehlo.multiply %12317, %12315 : tensor<1xf32> loc(#loc5236)
    %12319 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12320 = stablehlo.broadcast_in_dim %12319, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12321 = stablehlo.power %12320, %12306 : tensor<1xf32> loc(#loc5261)
    %12322 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12323 = stablehlo.broadcast_in_dim %12322, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12324 = stablehlo.subtract %12323, %12321 : tensor<1xf32> loc(#loc5260)
    %12325 = stablehlo.divide %12318, %12324 : tensor<1xf32> loc(#loc5253)
    %12326 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12327 = stablehlo.broadcast_in_dim %12326, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12328 = stablehlo.subtract %12327, %12302 : tensor<1xf32> loc(#loc5260)
    %12329 = stablehlo.broadcast_in_dim %12328, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %12330 = stablehlo.broadcast_in_dim %12329, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %12331 = stablehlo.multiply %12330, %4712 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %12332 = stablehlo.broadcast_in_dim %12302, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %12333 = stablehlo.broadcast_in_dim %12332, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %12334 = stablehlo.multiply %12333, %arg308 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %12335 = stablehlo.add %12331, %12334 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %12336 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12337 = stablehlo.broadcast_in_dim %12336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12338 = stablehlo.subtract %12337, %12325 : tensor<1xf32> loc(#loc5260)
    %12339 = stablehlo.multiply %4712, %4712 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %12340 = stablehlo.broadcast_in_dim %12338, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %12341 = stablehlo.broadcast_in_dim %12340, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %12342 = stablehlo.multiply %12341, %12339 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %12343 = stablehlo.broadcast_in_dim %12325, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %12344 = stablehlo.broadcast_in_dim %12343, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %12345 = stablehlo.multiply %12344, %arg484 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %12346 = stablehlo.add %12342, %12345 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %12347 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12348 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12349 = stablehlo.broadcast_in_dim %12348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12350 = stablehlo.add %12347, %12349 : tensor<1xf32> loc(#loc5259)
    %12351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12352 = stablehlo.broadcast_in_dim %12351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12353 = stablehlo.subtract %12350, %12352 : tensor<1xf32> loc(#loc5260)
    %12354 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12355 = stablehlo.broadcast_in_dim %12354, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12356 = stablehlo.power %12355, %12353 : tensor<1xf32> loc(#loc5261)
    %12357 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12358 = stablehlo.broadcast_in_dim %12357, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12359 = stablehlo.subtract %12358, %12356 : tensor<1xf32> loc(#loc5260)
    %12360 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12361 = stablehlo.broadcast_in_dim %12360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12362 = stablehlo.multiply %12361, %12359 : tensor<1xf32> loc(#loc5236)
    %12363 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12364 = stablehlo.broadcast_in_dim %12363, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12365 = stablehlo.power %12364, %12350 : tensor<1xf32> loc(#loc5261)
    %12366 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12367 = stablehlo.broadcast_in_dim %12366, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12368 = stablehlo.subtract %12367, %12365 : tensor<1xf32> loc(#loc5260)
    %12369 = stablehlo.divide %12362, %12368 : tensor<1xf32> loc(#loc5253)
    %12370 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12372 = stablehlo.broadcast_in_dim %12371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12373 = stablehlo.add %12370, %12372 : tensor<1xf32> loc(#loc5259)
    %12374 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12375 = stablehlo.broadcast_in_dim %12374, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12376 = stablehlo.subtract %12373, %12375 : tensor<1xf32> loc(#loc5260)
    %12377 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12378 = stablehlo.broadcast_in_dim %12377, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12379 = stablehlo.power %12378, %12376 : tensor<1xf32> loc(#loc5261)
    %12380 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12381 = stablehlo.broadcast_in_dim %12380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12382 = stablehlo.subtract %12381, %12379 : tensor<1xf32> loc(#loc5260)
    %12383 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12384 = stablehlo.broadcast_in_dim %12383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12385 = stablehlo.multiply %12384, %12382 : tensor<1xf32> loc(#loc5236)
    %12386 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12387 = stablehlo.broadcast_in_dim %12386, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12388 = stablehlo.power %12387, %12373 : tensor<1xf32> loc(#loc5261)
    %12389 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12390 = stablehlo.broadcast_in_dim %12389, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12391 = stablehlo.subtract %12390, %12388 : tensor<1xf32> loc(#loc5260)
    %12392 = stablehlo.divide %12385, %12391 : tensor<1xf32> loc(#loc5253)
    %12393 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12394 = stablehlo.broadcast_in_dim %12393, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12395 = stablehlo.subtract %12394, %12369 : tensor<1xf32> loc(#loc5260)
    %12396 = stablehlo.broadcast_in_dim %12395, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12397 = stablehlo.broadcast_in_dim %12396, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %12398 = stablehlo.multiply %12397, %4716 : tensor<1x8192xf32> loc(#loc5262)
    %12399 = stablehlo.broadcast_in_dim %12369, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12400 = stablehlo.broadcast_in_dim %12399, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %12401 = stablehlo.multiply %12400, %arg309 : tensor<1x8192xf32> loc(#loc5262)
    %12402 = stablehlo.add %12398, %12401 : tensor<1x8192xf32> loc(#loc5263)
    %12403 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12404 = stablehlo.broadcast_in_dim %12403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12405 = stablehlo.subtract %12404, %12392 : tensor<1xf32> loc(#loc5260)
    %12406 = stablehlo.multiply %4716, %4716 : tensor<1x8192xf32> loc(#loc5264)
    %12407 = stablehlo.broadcast_in_dim %12405, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12408 = stablehlo.broadcast_in_dim %12407, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %12409 = stablehlo.multiply %12408, %12406 : tensor<1x8192xf32> loc(#loc5262)
    %12410 = stablehlo.broadcast_in_dim %12392, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12411 = stablehlo.broadcast_in_dim %12410, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %12412 = stablehlo.multiply %12411, %arg485 : tensor<1x8192xf32> loc(#loc5262)
    %12413 = stablehlo.add %12409, %12412 : tensor<1x8192xf32> loc(#loc5263)
    %12414 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12415 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12416 = stablehlo.broadcast_in_dim %12415, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12417 = stablehlo.add %12414, %12416 : tensor<1xf32> loc(#loc5259)
    %12418 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12419 = stablehlo.broadcast_in_dim %12418, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12420 = stablehlo.subtract %12417, %12419 : tensor<1xf32> loc(#loc5260)
    %12421 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12422 = stablehlo.broadcast_in_dim %12421, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12423 = stablehlo.power %12422, %12420 : tensor<1xf32> loc(#loc5261)
    %12424 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12425 = stablehlo.broadcast_in_dim %12424, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12426 = stablehlo.subtract %12425, %12423 : tensor<1xf32> loc(#loc5260)
    %12427 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12428 = stablehlo.broadcast_in_dim %12427, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12429 = stablehlo.multiply %12428, %12426 : tensor<1xf32> loc(#loc5236)
    %12430 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12431 = stablehlo.broadcast_in_dim %12430, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12432 = stablehlo.power %12431, %12417 : tensor<1xf32> loc(#loc5261)
    %12433 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12434 = stablehlo.broadcast_in_dim %12433, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12435 = stablehlo.subtract %12434, %12432 : tensor<1xf32> loc(#loc5260)
    %12436 = stablehlo.divide %12429, %12435 : tensor<1xf32> loc(#loc5253)
    %12437 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12438 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12439 = stablehlo.broadcast_in_dim %12438, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12440 = stablehlo.add %12437, %12439 : tensor<1xf32> loc(#loc5259)
    %12441 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12442 = stablehlo.broadcast_in_dim %12441, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12443 = stablehlo.subtract %12440, %12442 : tensor<1xf32> loc(#loc5260)
    %12444 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12445 = stablehlo.broadcast_in_dim %12444, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12446 = stablehlo.power %12445, %12443 : tensor<1xf32> loc(#loc5261)
    %12447 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12448 = stablehlo.broadcast_in_dim %12447, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12449 = stablehlo.subtract %12448, %12446 : tensor<1xf32> loc(#loc5260)
    %12450 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12451 = stablehlo.broadcast_in_dim %12450, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12452 = stablehlo.multiply %12451, %12449 : tensor<1xf32> loc(#loc5236)
    %12453 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12454 = stablehlo.broadcast_in_dim %12453, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12455 = stablehlo.power %12454, %12440 : tensor<1xf32> loc(#loc5261)
    %12456 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12457 = stablehlo.broadcast_in_dim %12456, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12458 = stablehlo.subtract %12457, %12455 : tensor<1xf32> loc(#loc5260)
    %12459 = stablehlo.divide %12452, %12458 : tensor<1xf32> loc(#loc5253)
    %12460 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12461 = stablehlo.broadcast_in_dim %12460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12462 = stablehlo.subtract %12461, %12436 : tensor<1xf32> loc(#loc5260)
    %12463 = stablehlo.broadcast_in_dim %12462, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12464 = stablehlo.broadcast_in_dim %12463, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %12465 = stablehlo.multiply %12464, %4720 : tensor<1x2048x8192xf32> loc(#loc5262)
    %12466 = stablehlo.broadcast_in_dim %12436, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12467 = stablehlo.broadcast_in_dim %12466, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %12468 = stablehlo.multiply %12467, %arg310 : tensor<1x2048x8192xf32> loc(#loc5262)
    %12469 = stablehlo.add %12465, %12468 : tensor<1x2048x8192xf32> loc(#loc5263)
    %12470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12471 = stablehlo.broadcast_in_dim %12470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12472 = stablehlo.subtract %12471, %12459 : tensor<1xf32> loc(#loc5260)
    %12473 = stablehlo.multiply %4720, %4720 : tensor<1x2048x8192xf32> loc(#loc5264)
    %12474 = stablehlo.broadcast_in_dim %12472, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12475 = stablehlo.broadcast_in_dim %12474, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %12476 = stablehlo.multiply %12475, %12473 : tensor<1x2048x8192xf32> loc(#loc5262)
    %12477 = stablehlo.broadcast_in_dim %12459, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12478 = stablehlo.broadcast_in_dim %12477, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %12479 = stablehlo.multiply %12478, %arg486 : tensor<1x2048x8192xf32> loc(#loc5262)
    %12480 = stablehlo.add %12476, %12479 : tensor<1x2048x8192xf32> loc(#loc5263)
    %12481 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12482 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12483 = stablehlo.broadcast_in_dim %12482, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12484 = stablehlo.add %12481, %12483 : tensor<1xf32> loc(#loc5259)
    %12485 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12486 = stablehlo.broadcast_in_dim %12485, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12487 = stablehlo.subtract %12484, %12486 : tensor<1xf32> loc(#loc5260)
    %12488 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12489 = stablehlo.broadcast_in_dim %12488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12490 = stablehlo.power %12489, %12487 : tensor<1xf32> loc(#loc5261)
    %12491 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12492 = stablehlo.broadcast_in_dim %12491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12493 = stablehlo.subtract %12492, %12490 : tensor<1xf32> loc(#loc5260)
    %12494 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12495 = stablehlo.broadcast_in_dim %12494, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12496 = stablehlo.multiply %12495, %12493 : tensor<1xf32> loc(#loc5236)
    %12497 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12498 = stablehlo.broadcast_in_dim %12497, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12499 = stablehlo.power %12498, %12484 : tensor<1xf32> loc(#loc5261)
    %12500 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12501 = stablehlo.broadcast_in_dim %12500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12502 = stablehlo.subtract %12501, %12499 : tensor<1xf32> loc(#loc5260)
    %12503 = stablehlo.divide %12496, %12502 : tensor<1xf32> loc(#loc5253)
    %12504 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12505 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12506 = stablehlo.broadcast_in_dim %12505, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12507 = stablehlo.add %12504, %12506 : tensor<1xf32> loc(#loc5259)
    %12508 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12509 = stablehlo.broadcast_in_dim %12508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12510 = stablehlo.subtract %12507, %12509 : tensor<1xf32> loc(#loc5260)
    %12511 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12512 = stablehlo.broadcast_in_dim %12511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12513 = stablehlo.power %12512, %12510 : tensor<1xf32> loc(#loc5261)
    %12514 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12515 = stablehlo.broadcast_in_dim %12514, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12516 = stablehlo.subtract %12515, %12513 : tensor<1xf32> loc(#loc5260)
    %12517 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12518 = stablehlo.broadcast_in_dim %12517, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12519 = stablehlo.multiply %12518, %12516 : tensor<1xf32> loc(#loc5236)
    %12520 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12521 = stablehlo.broadcast_in_dim %12520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12522 = stablehlo.power %12521, %12507 : tensor<1xf32> loc(#loc5261)
    %12523 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12524 = stablehlo.broadcast_in_dim %12523, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12525 = stablehlo.subtract %12524, %12522 : tensor<1xf32> loc(#loc5260)
    %12526 = stablehlo.divide %12519, %12525 : tensor<1xf32> loc(#loc5253)
    %12527 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12528 = stablehlo.broadcast_in_dim %12527, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12529 = stablehlo.subtract %12528, %12503 : tensor<1xf32> loc(#loc5260)
    %12530 = stablehlo.broadcast_in_dim %12529, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12531 = stablehlo.broadcast_in_dim %12530, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12532 = stablehlo.multiply %12531, %4724 : tensor<1x2048xf32> loc(#loc5262)
    %12533 = stablehlo.broadcast_in_dim %12503, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12534 = stablehlo.broadcast_in_dim %12533, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12535 = stablehlo.multiply %12534, %arg311 : tensor<1x2048xf32> loc(#loc5262)
    %12536 = stablehlo.add %12532, %12535 : tensor<1x2048xf32> loc(#loc5263)
    %12537 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12538 = stablehlo.broadcast_in_dim %12537, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12539 = stablehlo.subtract %12538, %12526 : tensor<1xf32> loc(#loc5260)
    %12540 = stablehlo.multiply %4724, %4724 : tensor<1x2048xf32> loc(#loc5264)
    %12541 = stablehlo.broadcast_in_dim %12539, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12542 = stablehlo.broadcast_in_dim %12541, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12543 = stablehlo.multiply %12542, %12540 : tensor<1x2048xf32> loc(#loc5262)
    %12544 = stablehlo.broadcast_in_dim %12526, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12545 = stablehlo.broadcast_in_dim %12544, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12546 = stablehlo.multiply %12545, %arg487 : tensor<1x2048xf32> loc(#loc5262)
    %12547 = stablehlo.add %12543, %12546 : tensor<1x2048xf32> loc(#loc5263)
    %12548 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12549 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12550 = stablehlo.broadcast_in_dim %12549, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12551 = stablehlo.add %12548, %12550 : tensor<1xf32> loc(#loc5259)
    %12552 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12553 = stablehlo.broadcast_in_dim %12552, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12554 = stablehlo.subtract %12551, %12553 : tensor<1xf32> loc(#loc5260)
    %12555 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12556 = stablehlo.broadcast_in_dim %12555, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12557 = stablehlo.power %12556, %12554 : tensor<1xf32> loc(#loc5261)
    %12558 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12559 = stablehlo.broadcast_in_dim %12558, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12560 = stablehlo.subtract %12559, %12557 : tensor<1xf32> loc(#loc5260)
    %12561 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12562 = stablehlo.broadcast_in_dim %12561, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12563 = stablehlo.multiply %12562, %12560 : tensor<1xf32> loc(#loc5236)
    %12564 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12565 = stablehlo.broadcast_in_dim %12564, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12566 = stablehlo.power %12565, %12551 : tensor<1xf32> loc(#loc5261)
    %12567 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12568 = stablehlo.broadcast_in_dim %12567, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12569 = stablehlo.subtract %12568, %12566 : tensor<1xf32> loc(#loc5260)
    %12570 = stablehlo.divide %12563, %12569 : tensor<1xf32> loc(#loc5253)
    %12571 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12572 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12573 = stablehlo.broadcast_in_dim %12572, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12574 = stablehlo.add %12571, %12573 : tensor<1xf32> loc(#loc5259)
    %12575 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12576 = stablehlo.broadcast_in_dim %12575, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12577 = stablehlo.subtract %12574, %12576 : tensor<1xf32> loc(#loc5260)
    %12578 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12579 = stablehlo.broadcast_in_dim %12578, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12580 = stablehlo.power %12579, %12577 : tensor<1xf32> loc(#loc5261)
    %12581 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12582 = stablehlo.broadcast_in_dim %12581, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12583 = stablehlo.subtract %12582, %12580 : tensor<1xf32> loc(#loc5260)
    %12584 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12585 = stablehlo.broadcast_in_dim %12584, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12586 = stablehlo.multiply %12585, %12583 : tensor<1xf32> loc(#loc5236)
    %12587 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12588 = stablehlo.broadcast_in_dim %12587, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12589 = stablehlo.power %12588, %12574 : tensor<1xf32> loc(#loc5261)
    %12590 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12591 = stablehlo.broadcast_in_dim %12590, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12592 = stablehlo.subtract %12591, %12589 : tensor<1xf32> loc(#loc5260)
    %12593 = stablehlo.divide %12586, %12592 : tensor<1xf32> loc(#loc5253)
    %12594 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12595 = stablehlo.broadcast_in_dim %12594, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12596 = stablehlo.subtract %12595, %12570 : tensor<1xf32> loc(#loc5260)
    %12597 = stablehlo.broadcast_in_dim %12596, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12598 = stablehlo.broadcast_in_dim %12597, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %12599 = stablehlo.multiply %12598, %4728 : tensor<1x8192x2048xf32> loc(#loc5262)
    %12600 = stablehlo.broadcast_in_dim %12570, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12601 = stablehlo.broadcast_in_dim %12600, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %12602 = stablehlo.multiply %12601, %arg312 : tensor<1x8192x2048xf32> loc(#loc5262)
    %12603 = stablehlo.add %12599, %12602 : tensor<1x8192x2048xf32> loc(#loc5263)
    %12604 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12605 = stablehlo.broadcast_in_dim %12604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12606 = stablehlo.subtract %12605, %12593 : tensor<1xf32> loc(#loc5260)
    %12607 = stablehlo.multiply %4728, %4728 : tensor<1x8192x2048xf32> loc(#loc5264)
    %12608 = stablehlo.broadcast_in_dim %12606, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12609 = stablehlo.broadcast_in_dim %12608, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %12610 = stablehlo.multiply %12609, %12607 : tensor<1x8192x2048xf32> loc(#loc5262)
    %12611 = stablehlo.broadcast_in_dim %12593, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %12612 = stablehlo.broadcast_in_dim %12611, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %12613 = stablehlo.multiply %12612, %arg488 : tensor<1x8192x2048xf32> loc(#loc5262)
    %12614 = stablehlo.add %12610, %12613 : tensor<1x8192x2048xf32> loc(#loc5263)
    %12615 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12616 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12617 = stablehlo.broadcast_in_dim %12616, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12618 = stablehlo.add %12615, %12617 : tensor<1xf32> loc(#loc5259)
    %12619 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12620 = stablehlo.broadcast_in_dim %12619, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12621 = stablehlo.subtract %12618, %12620 : tensor<1xf32> loc(#loc5260)
    %12622 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12623 = stablehlo.broadcast_in_dim %12622, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12624 = stablehlo.power %12623, %12621 : tensor<1xf32> loc(#loc5261)
    %12625 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12626 = stablehlo.broadcast_in_dim %12625, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12627 = stablehlo.subtract %12626, %12624 : tensor<1xf32> loc(#loc5260)
    %12628 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12629 = stablehlo.broadcast_in_dim %12628, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12630 = stablehlo.multiply %12629, %12627 : tensor<1xf32> loc(#loc5236)
    %12631 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12632 = stablehlo.broadcast_in_dim %12631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12633 = stablehlo.power %12632, %12618 : tensor<1xf32> loc(#loc5261)
    %12634 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12635 = stablehlo.broadcast_in_dim %12634, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12636 = stablehlo.subtract %12635, %12633 : tensor<1xf32> loc(#loc5260)
    %12637 = stablehlo.divide %12630, %12636 : tensor<1xf32> loc(#loc5253)
    %12638 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12639 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12640 = stablehlo.broadcast_in_dim %12639, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12641 = stablehlo.add %12638, %12640 : tensor<1xf32> loc(#loc5259)
    %12642 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12643 = stablehlo.broadcast_in_dim %12642, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12644 = stablehlo.subtract %12641, %12643 : tensor<1xf32> loc(#loc5260)
    %12645 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12646 = stablehlo.broadcast_in_dim %12645, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12647 = stablehlo.power %12646, %12644 : tensor<1xf32> loc(#loc5261)
    %12648 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12649 = stablehlo.broadcast_in_dim %12648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12650 = stablehlo.subtract %12649, %12647 : tensor<1xf32> loc(#loc5260)
    %12651 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12652 = stablehlo.broadcast_in_dim %12651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12653 = stablehlo.multiply %12652, %12650 : tensor<1xf32> loc(#loc5236)
    %12654 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12655 = stablehlo.broadcast_in_dim %12654, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12656 = stablehlo.power %12655, %12641 : tensor<1xf32> loc(#loc5261)
    %12657 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12658 = stablehlo.broadcast_in_dim %12657, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12659 = stablehlo.subtract %12658, %12656 : tensor<1xf32> loc(#loc5260)
    %12660 = stablehlo.divide %12653, %12659 : tensor<1xf32> loc(#loc5253)
    %12661 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12662 = stablehlo.broadcast_in_dim %12661, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12663 = stablehlo.subtract %12662, %12637 : tensor<1xf32> loc(#loc5260)
    %12664 = stablehlo.broadcast_in_dim %12663, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12665 = stablehlo.broadcast_in_dim %12664, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12666 = stablehlo.multiply %12665, %4732 : tensor<1x2048xf32> loc(#loc5262)
    %12667 = stablehlo.broadcast_in_dim %12637, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12668 = stablehlo.broadcast_in_dim %12667, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12669 = stablehlo.multiply %12668, %arg313 : tensor<1x2048xf32> loc(#loc5262)
    %12670 = stablehlo.add %12666, %12669 : tensor<1x2048xf32> loc(#loc5263)
    %12671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12672 = stablehlo.broadcast_in_dim %12671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12673 = stablehlo.subtract %12672, %12660 : tensor<1xf32> loc(#loc5260)
    %12674 = stablehlo.multiply %4732, %4732 : tensor<1x2048xf32> loc(#loc5264)
    %12675 = stablehlo.broadcast_in_dim %12673, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12676 = stablehlo.broadcast_in_dim %12675, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12677 = stablehlo.multiply %12676, %12674 : tensor<1x2048xf32> loc(#loc5262)
    %12678 = stablehlo.broadcast_in_dim %12660, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12679 = stablehlo.broadcast_in_dim %12678, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12680 = stablehlo.multiply %12679, %arg489 : tensor<1x2048xf32> loc(#loc5262)
    %12681 = stablehlo.add %12677, %12680 : tensor<1x2048xf32> loc(#loc5263)
    %12682 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12683 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12684 = stablehlo.broadcast_in_dim %12683, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12685 = stablehlo.add %12682, %12684 : tensor<1xf32> loc(#loc5259)
    %12686 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12687 = stablehlo.broadcast_in_dim %12686, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12688 = stablehlo.subtract %12685, %12687 : tensor<1xf32> loc(#loc5260)
    %12689 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12690 = stablehlo.broadcast_in_dim %12689, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12691 = stablehlo.power %12690, %12688 : tensor<1xf32> loc(#loc5261)
    %12692 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12693 = stablehlo.broadcast_in_dim %12692, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12694 = stablehlo.subtract %12693, %12691 : tensor<1xf32> loc(#loc5260)
    %12695 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12696 = stablehlo.broadcast_in_dim %12695, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12697 = stablehlo.multiply %12696, %12694 : tensor<1xf32> loc(#loc5236)
    %12698 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12699 = stablehlo.broadcast_in_dim %12698, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12700 = stablehlo.power %12699, %12685 : tensor<1xf32> loc(#loc5261)
    %12701 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12702 = stablehlo.broadcast_in_dim %12701, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12703 = stablehlo.subtract %12702, %12700 : tensor<1xf32> loc(#loc5260)
    %12704 = stablehlo.divide %12697, %12703 : tensor<1xf32> loc(#loc5253)
    %12705 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12706 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12707 = stablehlo.broadcast_in_dim %12706, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12708 = stablehlo.add %12705, %12707 : tensor<1xf32> loc(#loc5259)
    %12709 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12710 = stablehlo.broadcast_in_dim %12709, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12711 = stablehlo.subtract %12708, %12710 : tensor<1xf32> loc(#loc5260)
    %12712 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12713 = stablehlo.broadcast_in_dim %12712, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12714 = stablehlo.power %12713, %12711 : tensor<1xf32> loc(#loc5261)
    %12715 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12716 = stablehlo.broadcast_in_dim %12715, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12717 = stablehlo.subtract %12716, %12714 : tensor<1xf32> loc(#loc5260)
    %12718 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12719 = stablehlo.broadcast_in_dim %12718, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12720 = stablehlo.multiply %12719, %12717 : tensor<1xf32> loc(#loc5236)
    %12721 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12722 = stablehlo.broadcast_in_dim %12721, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12723 = stablehlo.power %12722, %12708 : tensor<1xf32> loc(#loc5261)
    %12724 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12725 = stablehlo.broadcast_in_dim %12724, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12726 = stablehlo.subtract %12725, %12723 : tensor<1xf32> loc(#loc5260)
    %12727 = stablehlo.divide %12720, %12726 : tensor<1xf32> loc(#loc5253)
    %12728 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12729 = stablehlo.broadcast_in_dim %12728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12730 = stablehlo.subtract %12729, %12704 : tensor<1xf32> loc(#loc5260)
    %12731 = stablehlo.broadcast_in_dim %12730, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12732 = stablehlo.broadcast_in_dim %12731, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12733 = stablehlo.multiply %12732, %4736 : tensor<1x2048xf32> loc(#loc5262)
    %12734 = stablehlo.broadcast_in_dim %12704, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12735 = stablehlo.broadcast_in_dim %12734, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12736 = stablehlo.multiply %12735, %arg314 : tensor<1x2048xf32> loc(#loc5262)
    %12737 = stablehlo.add %12733, %12736 : tensor<1x2048xf32> loc(#loc5263)
    %12738 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12739 = stablehlo.broadcast_in_dim %12738, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12740 = stablehlo.subtract %12739, %12727 : tensor<1xf32> loc(#loc5260)
    %12741 = stablehlo.multiply %4736, %4736 : tensor<1x2048xf32> loc(#loc5264)
    %12742 = stablehlo.broadcast_in_dim %12740, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12743 = stablehlo.broadcast_in_dim %12742, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12744 = stablehlo.multiply %12743, %12741 : tensor<1x2048xf32> loc(#loc5262)
    %12745 = stablehlo.broadcast_in_dim %12727, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12746 = stablehlo.broadcast_in_dim %12745, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12747 = stablehlo.multiply %12746, %arg490 : tensor<1x2048xf32> loc(#loc5262)
    %12748 = stablehlo.add %12744, %12747 : tensor<1x2048xf32> loc(#loc5263)
    %12749 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12750 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12751 = stablehlo.broadcast_in_dim %12750, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12752 = stablehlo.add %12749, %12751 : tensor<1xf32> loc(#loc5259)
    %12753 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12754 = stablehlo.broadcast_in_dim %12753, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12755 = stablehlo.subtract %12752, %12754 : tensor<1xf32> loc(#loc5260)
    %12756 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12757 = stablehlo.broadcast_in_dim %12756, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12758 = stablehlo.power %12757, %12755 : tensor<1xf32> loc(#loc5261)
    %12759 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12760 = stablehlo.broadcast_in_dim %12759, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12761 = stablehlo.subtract %12760, %12758 : tensor<1xf32> loc(#loc5260)
    %12762 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12763 = stablehlo.broadcast_in_dim %12762, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12764 = stablehlo.multiply %12763, %12761 : tensor<1xf32> loc(#loc5236)
    %12765 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12766 = stablehlo.broadcast_in_dim %12765, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12767 = stablehlo.power %12766, %12752 : tensor<1xf32> loc(#loc5261)
    %12768 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12769 = stablehlo.broadcast_in_dim %12768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12770 = stablehlo.subtract %12769, %12767 : tensor<1xf32> loc(#loc5260)
    %12771 = stablehlo.divide %12764, %12770 : tensor<1xf32> loc(#loc5253)
    %12772 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12773 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12774 = stablehlo.broadcast_in_dim %12773, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12775 = stablehlo.add %12772, %12774 : tensor<1xf32> loc(#loc5259)
    %12776 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12777 = stablehlo.broadcast_in_dim %12776, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12778 = stablehlo.subtract %12775, %12777 : tensor<1xf32> loc(#loc5260)
    %12779 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12780 = stablehlo.broadcast_in_dim %12779, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12781 = stablehlo.power %12780, %12778 : tensor<1xf32> loc(#loc5261)
    %12782 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12783 = stablehlo.broadcast_in_dim %12782, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12784 = stablehlo.subtract %12783, %12781 : tensor<1xf32> loc(#loc5260)
    %12785 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12786 = stablehlo.broadcast_in_dim %12785, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12787 = stablehlo.multiply %12786, %12784 : tensor<1xf32> loc(#loc5236)
    %12788 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12789 = stablehlo.broadcast_in_dim %12788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12790 = stablehlo.power %12789, %12775 : tensor<1xf32> loc(#loc5261)
    %12791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12792 = stablehlo.broadcast_in_dim %12791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12793 = stablehlo.subtract %12792, %12790 : tensor<1xf32> loc(#loc5260)
    %12794 = stablehlo.divide %12787, %12793 : tensor<1xf32> loc(#loc5253)
    %12795 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12796 = stablehlo.broadcast_in_dim %12795, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12797 = stablehlo.subtract %12796, %12771 : tensor<1xf32> loc(#loc5260)
    %12798 = stablehlo.broadcast_in_dim %12797, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12799 = stablehlo.broadcast_in_dim %12798, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12800 = stablehlo.multiply %12799, %4740 : tensor<1x2048xf32> loc(#loc5262)
    %12801 = stablehlo.broadcast_in_dim %12771, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12802 = stablehlo.broadcast_in_dim %12801, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12803 = stablehlo.multiply %12802, %arg315 : tensor<1x2048xf32> loc(#loc5262)
    %12804 = stablehlo.add %12800, %12803 : tensor<1x2048xf32> loc(#loc5263)
    %12805 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12806 = stablehlo.broadcast_in_dim %12805, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12807 = stablehlo.subtract %12806, %12794 : tensor<1xf32> loc(#loc5260)
    %12808 = stablehlo.multiply %4740, %4740 : tensor<1x2048xf32> loc(#loc5264)
    %12809 = stablehlo.broadcast_in_dim %12807, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12810 = stablehlo.broadcast_in_dim %12809, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12811 = stablehlo.multiply %12810, %12808 : tensor<1x2048xf32> loc(#loc5262)
    %12812 = stablehlo.broadcast_in_dim %12794, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12813 = stablehlo.broadcast_in_dim %12812, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12814 = stablehlo.multiply %12813, %arg491 : tensor<1x2048xf32> loc(#loc5262)
    %12815 = stablehlo.add %12811, %12814 : tensor<1x2048xf32> loc(#loc5263)
    %12816 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12817 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12818 = stablehlo.broadcast_in_dim %12817, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12819 = stablehlo.add %12816, %12818 : tensor<1xf32> loc(#loc5259)
    %12820 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12821 = stablehlo.broadcast_in_dim %12820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12822 = stablehlo.subtract %12819, %12821 : tensor<1xf32> loc(#loc5260)
    %12823 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12824 = stablehlo.broadcast_in_dim %12823, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12825 = stablehlo.power %12824, %12822 : tensor<1xf32> loc(#loc5261)
    %12826 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12827 = stablehlo.broadcast_in_dim %12826, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12828 = stablehlo.subtract %12827, %12825 : tensor<1xf32> loc(#loc5260)
    %12829 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12830 = stablehlo.broadcast_in_dim %12829, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12831 = stablehlo.multiply %12830, %12828 : tensor<1xf32> loc(#loc5236)
    %12832 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12833 = stablehlo.broadcast_in_dim %12832, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12834 = stablehlo.power %12833, %12819 : tensor<1xf32> loc(#loc5261)
    %12835 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12836 = stablehlo.broadcast_in_dim %12835, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12837 = stablehlo.subtract %12836, %12834 : tensor<1xf32> loc(#loc5260)
    %12838 = stablehlo.divide %12831, %12837 : tensor<1xf32> loc(#loc5253)
    %12839 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12840 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12841 = stablehlo.broadcast_in_dim %12840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12842 = stablehlo.add %12839, %12841 : tensor<1xf32> loc(#loc5259)
    %12843 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12844 = stablehlo.broadcast_in_dim %12843, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12845 = stablehlo.subtract %12842, %12844 : tensor<1xf32> loc(#loc5260)
    %12846 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12847 = stablehlo.broadcast_in_dim %12846, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12848 = stablehlo.power %12847, %12845 : tensor<1xf32> loc(#loc5261)
    %12849 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12850 = stablehlo.broadcast_in_dim %12849, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12851 = stablehlo.subtract %12850, %12848 : tensor<1xf32> loc(#loc5260)
    %12852 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12853 = stablehlo.broadcast_in_dim %12852, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12854 = stablehlo.multiply %12853, %12851 : tensor<1xf32> loc(#loc5236)
    %12855 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12856 = stablehlo.broadcast_in_dim %12855, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12857 = stablehlo.power %12856, %12842 : tensor<1xf32> loc(#loc5261)
    %12858 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12859 = stablehlo.broadcast_in_dim %12858, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12860 = stablehlo.subtract %12859, %12857 : tensor<1xf32> loc(#loc5260)
    %12861 = stablehlo.divide %12854, %12860 : tensor<1xf32> loc(#loc5253)
    %12862 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12863 = stablehlo.broadcast_in_dim %12862, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12864 = stablehlo.subtract %12863, %12838 : tensor<1xf32> loc(#loc5260)
    %12865 = stablehlo.broadcast_in_dim %12864, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12866 = stablehlo.broadcast_in_dim %12865, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12867 = stablehlo.multiply %12866, %4744 : tensor<1x2048xf32> loc(#loc5262)
    %12868 = stablehlo.broadcast_in_dim %12838, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12869 = stablehlo.broadcast_in_dim %12868, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12870 = stablehlo.multiply %12869, %arg316 : tensor<1x2048xf32> loc(#loc5262)
    %12871 = stablehlo.add %12867, %12870 : tensor<1x2048xf32> loc(#loc5263)
    %12872 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12873 = stablehlo.broadcast_in_dim %12872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12874 = stablehlo.subtract %12873, %12861 : tensor<1xf32> loc(#loc5260)
    %12875 = stablehlo.multiply %4744, %4744 : tensor<1x2048xf32> loc(#loc5264)
    %12876 = stablehlo.broadcast_in_dim %12874, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12877 = stablehlo.broadcast_in_dim %12876, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12878 = stablehlo.multiply %12877, %12875 : tensor<1x2048xf32> loc(#loc5262)
    %12879 = stablehlo.broadcast_in_dim %12861, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %12880 = stablehlo.broadcast_in_dim %12879, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %12881 = stablehlo.multiply %12880, %arg492 : tensor<1x2048xf32> loc(#loc5262)
    %12882 = stablehlo.add %12878, %12881 : tensor<1x2048xf32> loc(#loc5263)
    %12883 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12884 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12885 = stablehlo.broadcast_in_dim %12884, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12886 = stablehlo.add %12883, %12885 : tensor<1xf32> loc(#loc5259)
    %12887 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12888 = stablehlo.broadcast_in_dim %12887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12889 = stablehlo.subtract %12886, %12888 : tensor<1xf32> loc(#loc5260)
    %12890 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12891 = stablehlo.broadcast_in_dim %12890, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12892 = stablehlo.power %12891, %12889 : tensor<1xf32> loc(#loc5261)
    %12893 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12894 = stablehlo.broadcast_in_dim %12893, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12895 = stablehlo.subtract %12894, %12892 : tensor<1xf32> loc(#loc5260)
    %12896 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12897 = stablehlo.broadcast_in_dim %12896, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12898 = stablehlo.multiply %12897, %12895 : tensor<1xf32> loc(#loc5236)
    %12899 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12900 = stablehlo.broadcast_in_dim %12899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12901 = stablehlo.power %12900, %12886 : tensor<1xf32> loc(#loc5261)
    %12902 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12903 = stablehlo.broadcast_in_dim %12902, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12904 = stablehlo.subtract %12903, %12901 : tensor<1xf32> loc(#loc5260)
    %12905 = stablehlo.divide %12898, %12904 : tensor<1xf32> loc(#loc5253)
    %12906 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12907 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12908 = stablehlo.broadcast_in_dim %12907, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12909 = stablehlo.add %12906, %12908 : tensor<1xf32> loc(#loc5259)
    %12910 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12911 = stablehlo.broadcast_in_dim %12910, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12912 = stablehlo.subtract %12909, %12911 : tensor<1xf32> loc(#loc5260)
    %12913 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12914 = stablehlo.broadcast_in_dim %12913, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12915 = stablehlo.power %12914, %12912 : tensor<1xf32> loc(#loc5261)
    %12916 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12917 = stablehlo.broadcast_in_dim %12916, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12918 = stablehlo.subtract %12917, %12915 : tensor<1xf32> loc(#loc5260)
    %12919 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12920 = stablehlo.broadcast_in_dim %12919, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12921 = stablehlo.multiply %12920, %12918 : tensor<1xf32> loc(#loc5236)
    %12922 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12923 = stablehlo.broadcast_in_dim %12922, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12924 = stablehlo.power %12923, %12909 : tensor<1xf32> loc(#loc5261)
    %12925 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12926 = stablehlo.broadcast_in_dim %12925, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12927 = stablehlo.subtract %12926, %12924 : tensor<1xf32> loc(#loc5260)
    %12928 = stablehlo.divide %12921, %12927 : tensor<1xf32> loc(#loc5253)
    %12929 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12930 = stablehlo.broadcast_in_dim %12929, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12931 = stablehlo.subtract %12930, %12905 : tensor<1xf32> loc(#loc5260)
    %12932 = stablehlo.broadcast_in_dim %12931, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12933 = stablehlo.broadcast_in_dim %12932, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12934 = stablehlo.multiply %12933, %4748 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12935 = stablehlo.broadcast_in_dim %12905, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12936 = stablehlo.broadcast_in_dim %12935, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12937 = stablehlo.multiply %12936, %arg317 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12938 = stablehlo.add %12934, %12937 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %12939 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12940 = stablehlo.broadcast_in_dim %12939, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12941 = stablehlo.subtract %12940, %12928 : tensor<1xf32> loc(#loc5260)
    %12942 = stablehlo.multiply %4748, %4748 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %12943 = stablehlo.broadcast_in_dim %12941, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12944 = stablehlo.broadcast_in_dim %12943, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12945 = stablehlo.multiply %12944, %12942 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12946 = stablehlo.broadcast_in_dim %12928, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %12947 = stablehlo.broadcast_in_dim %12946, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12948 = stablehlo.multiply %12947, %arg493 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %12949 = stablehlo.add %12945, %12948 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %12950 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12951 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12952 = stablehlo.broadcast_in_dim %12951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12953 = stablehlo.add %12950, %12952 : tensor<1xf32> loc(#loc5259)
    %12954 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12955 = stablehlo.broadcast_in_dim %12954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12956 = stablehlo.subtract %12953, %12955 : tensor<1xf32> loc(#loc5260)
    %12957 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12958 = stablehlo.broadcast_in_dim %12957, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12959 = stablehlo.power %12958, %12956 : tensor<1xf32> loc(#loc5261)
    %12960 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12961 = stablehlo.broadcast_in_dim %12960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12962 = stablehlo.subtract %12961, %12959 : tensor<1xf32> loc(#loc5260)
    %12963 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12964 = stablehlo.broadcast_in_dim %12963, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12965 = stablehlo.multiply %12964, %12962 : tensor<1xf32> loc(#loc5236)
    %12966 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %12967 = stablehlo.broadcast_in_dim %12966, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12968 = stablehlo.power %12967, %12953 : tensor<1xf32> loc(#loc5261)
    %12969 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12970 = stablehlo.broadcast_in_dim %12969, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12971 = stablehlo.subtract %12970, %12968 : tensor<1xf32> loc(#loc5260)
    %12972 = stablehlo.divide %12965, %12971 : tensor<1xf32> loc(#loc5253)
    %12973 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %12974 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12975 = stablehlo.broadcast_in_dim %12974, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %12976 = stablehlo.add %12973, %12975 : tensor<1xf32> loc(#loc5259)
    %12977 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12978 = stablehlo.broadcast_in_dim %12977, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12979 = stablehlo.subtract %12976, %12978 : tensor<1xf32> loc(#loc5260)
    %12980 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12981 = stablehlo.broadcast_in_dim %12980, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12982 = stablehlo.power %12981, %12979 : tensor<1xf32> loc(#loc5261)
    %12983 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12984 = stablehlo.broadcast_in_dim %12983, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12985 = stablehlo.subtract %12984, %12982 : tensor<1xf32> loc(#loc5260)
    %12986 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12987 = stablehlo.broadcast_in_dim %12986, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %12988 = stablehlo.multiply %12987, %12985 : tensor<1xf32> loc(#loc5236)
    %12989 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %12990 = stablehlo.broadcast_in_dim %12989, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %12991 = stablehlo.power %12990, %12976 : tensor<1xf32> loc(#loc5261)
    %12992 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12993 = stablehlo.broadcast_in_dim %12992, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12994 = stablehlo.subtract %12993, %12991 : tensor<1xf32> loc(#loc5260)
    %12995 = stablehlo.divide %12988, %12994 : tensor<1xf32> loc(#loc5253)
    %12996 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %12997 = stablehlo.broadcast_in_dim %12996, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %12998 = stablehlo.subtract %12997, %12972 : tensor<1xf32> loc(#loc5260)
    %12999 = stablehlo.broadcast_in_dim %12998, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13000 = stablehlo.broadcast_in_dim %12999, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13001 = stablehlo.multiply %13000, %4752 : tensor<1x64xf32> loc(#loc5262)
    %13002 = stablehlo.broadcast_in_dim %12972, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13003 = stablehlo.broadcast_in_dim %13002, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13004 = stablehlo.multiply %13003, %arg318 : tensor<1x64xf32> loc(#loc5262)
    %13005 = stablehlo.add %13001, %13004 : tensor<1x64xf32> loc(#loc5263)
    %13006 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13007 = stablehlo.broadcast_in_dim %13006, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13008 = stablehlo.subtract %13007, %12995 : tensor<1xf32> loc(#loc5260)
    %13009 = stablehlo.multiply %4752, %4752 : tensor<1x64xf32> loc(#loc5264)
    %13010 = stablehlo.broadcast_in_dim %13008, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13011 = stablehlo.broadcast_in_dim %13010, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13012 = stablehlo.multiply %13011, %13009 : tensor<1x64xf32> loc(#loc5262)
    %13013 = stablehlo.broadcast_in_dim %12995, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13014 = stablehlo.broadcast_in_dim %13013, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13015 = stablehlo.multiply %13014, %arg494 : tensor<1x64xf32> loc(#loc5262)
    %13016 = stablehlo.add %13012, %13015 : tensor<1x64xf32> loc(#loc5263)
    %13017 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13018 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13019 = stablehlo.broadcast_in_dim %13018, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13020 = stablehlo.add %13017, %13019 : tensor<1xf32> loc(#loc5259)
    %13021 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13022 = stablehlo.broadcast_in_dim %13021, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13023 = stablehlo.subtract %13020, %13022 : tensor<1xf32> loc(#loc5260)
    %13024 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13025 = stablehlo.broadcast_in_dim %13024, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13026 = stablehlo.power %13025, %13023 : tensor<1xf32> loc(#loc5261)
    %13027 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13028 = stablehlo.broadcast_in_dim %13027, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13029 = stablehlo.subtract %13028, %13026 : tensor<1xf32> loc(#loc5260)
    %13030 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13031 = stablehlo.broadcast_in_dim %13030, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13032 = stablehlo.multiply %13031, %13029 : tensor<1xf32> loc(#loc5236)
    %13033 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13034 = stablehlo.broadcast_in_dim %13033, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13035 = stablehlo.power %13034, %13020 : tensor<1xf32> loc(#loc5261)
    %13036 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13037 = stablehlo.broadcast_in_dim %13036, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13038 = stablehlo.subtract %13037, %13035 : tensor<1xf32> loc(#loc5260)
    %13039 = stablehlo.divide %13032, %13038 : tensor<1xf32> loc(#loc5253)
    %13040 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13041 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13042 = stablehlo.broadcast_in_dim %13041, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13043 = stablehlo.add %13040, %13042 : tensor<1xf32> loc(#loc5259)
    %13044 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13045 = stablehlo.broadcast_in_dim %13044, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13046 = stablehlo.subtract %13043, %13045 : tensor<1xf32> loc(#loc5260)
    %13047 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13048 = stablehlo.broadcast_in_dim %13047, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13049 = stablehlo.power %13048, %13046 : tensor<1xf32> loc(#loc5261)
    %13050 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13051 = stablehlo.broadcast_in_dim %13050, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13052 = stablehlo.subtract %13051, %13049 : tensor<1xf32> loc(#loc5260)
    %13053 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13054 = stablehlo.broadcast_in_dim %13053, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13055 = stablehlo.multiply %13054, %13052 : tensor<1xf32> loc(#loc5236)
    %13056 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13057 = stablehlo.broadcast_in_dim %13056, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13058 = stablehlo.power %13057, %13043 : tensor<1xf32> loc(#loc5261)
    %13059 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13060 = stablehlo.broadcast_in_dim %13059, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13061 = stablehlo.subtract %13060, %13058 : tensor<1xf32> loc(#loc5260)
    %13062 = stablehlo.divide %13055, %13061 : tensor<1xf32> loc(#loc5253)
    %13063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13064 = stablehlo.broadcast_in_dim %13063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13065 = stablehlo.subtract %13064, %13039 : tensor<1xf32> loc(#loc5260)
    %13066 = stablehlo.broadcast_in_dim %13065, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13067 = stablehlo.broadcast_in_dim %13066, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13068 = stablehlo.multiply %13067, %4756 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13069 = stablehlo.broadcast_in_dim %13039, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13070 = stablehlo.broadcast_in_dim %13069, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13071 = stablehlo.multiply %13070, %arg319 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13072 = stablehlo.add %13068, %13071 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %13073 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13074 = stablehlo.broadcast_in_dim %13073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13075 = stablehlo.subtract %13074, %13062 : tensor<1xf32> loc(#loc5260)
    %13076 = stablehlo.multiply %4756, %4756 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %13077 = stablehlo.broadcast_in_dim %13075, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13078 = stablehlo.broadcast_in_dim %13077, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13079 = stablehlo.multiply %13078, %13076 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13080 = stablehlo.broadcast_in_dim %13062, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13081 = stablehlo.broadcast_in_dim %13080, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13082 = stablehlo.multiply %13081, %arg495 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13083 = stablehlo.add %13079, %13082 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %13084 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13085 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13086 = stablehlo.broadcast_in_dim %13085, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13087 = stablehlo.add %13084, %13086 : tensor<1xf32> loc(#loc5259)
    %13088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13089 = stablehlo.broadcast_in_dim %13088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13090 = stablehlo.subtract %13087, %13089 : tensor<1xf32> loc(#loc5260)
    %13091 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13092 = stablehlo.broadcast_in_dim %13091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13093 = stablehlo.power %13092, %13090 : tensor<1xf32> loc(#loc5261)
    %13094 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13095 = stablehlo.broadcast_in_dim %13094, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13096 = stablehlo.subtract %13095, %13093 : tensor<1xf32> loc(#loc5260)
    %13097 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13098 = stablehlo.broadcast_in_dim %13097, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13099 = stablehlo.multiply %13098, %13096 : tensor<1xf32> loc(#loc5236)
    %13100 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13101 = stablehlo.broadcast_in_dim %13100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13102 = stablehlo.power %13101, %13087 : tensor<1xf32> loc(#loc5261)
    %13103 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13104 = stablehlo.broadcast_in_dim %13103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13105 = stablehlo.subtract %13104, %13102 : tensor<1xf32> loc(#loc5260)
    %13106 = stablehlo.divide %13099, %13105 : tensor<1xf32> loc(#loc5253)
    %13107 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13108 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13109 = stablehlo.broadcast_in_dim %13108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13110 = stablehlo.add %13107, %13109 : tensor<1xf32> loc(#loc5259)
    %13111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13112 = stablehlo.broadcast_in_dim %13111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13113 = stablehlo.subtract %13110, %13112 : tensor<1xf32> loc(#loc5260)
    %13114 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13115 = stablehlo.broadcast_in_dim %13114, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13116 = stablehlo.power %13115, %13113 : tensor<1xf32> loc(#loc5261)
    %13117 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13118 = stablehlo.broadcast_in_dim %13117, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13119 = stablehlo.subtract %13118, %13116 : tensor<1xf32> loc(#loc5260)
    %13120 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13121 = stablehlo.broadcast_in_dim %13120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13122 = stablehlo.multiply %13121, %13119 : tensor<1xf32> loc(#loc5236)
    %13123 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13124 = stablehlo.broadcast_in_dim %13123, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13125 = stablehlo.power %13124, %13110 : tensor<1xf32> loc(#loc5261)
    %13126 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13127 = stablehlo.broadcast_in_dim %13126, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13128 = stablehlo.subtract %13127, %13125 : tensor<1xf32> loc(#loc5260)
    %13129 = stablehlo.divide %13122, %13128 : tensor<1xf32> loc(#loc5253)
    %13130 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13131 = stablehlo.broadcast_in_dim %13130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13132 = stablehlo.subtract %13131, %13106 : tensor<1xf32> loc(#loc5260)
    %13133 = stablehlo.broadcast_in_dim %13132, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13134 = stablehlo.broadcast_in_dim %13133, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13135 = stablehlo.multiply %13134, %4760 : tensor<1x8192xf32> loc(#loc5262)
    %13136 = stablehlo.broadcast_in_dim %13106, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13137 = stablehlo.broadcast_in_dim %13136, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13138 = stablehlo.multiply %13137, %arg320 : tensor<1x8192xf32> loc(#loc5262)
    %13139 = stablehlo.add %13135, %13138 : tensor<1x8192xf32> loc(#loc5263)
    %13140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13141 = stablehlo.broadcast_in_dim %13140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13142 = stablehlo.subtract %13141, %13129 : tensor<1xf32> loc(#loc5260)
    %13143 = stablehlo.multiply %4760, %4760 : tensor<1x8192xf32> loc(#loc5264)
    %13144 = stablehlo.broadcast_in_dim %13142, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13145 = stablehlo.broadcast_in_dim %13144, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13146 = stablehlo.multiply %13145, %13143 : tensor<1x8192xf32> loc(#loc5262)
    %13147 = stablehlo.broadcast_in_dim %13129, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13148 = stablehlo.broadcast_in_dim %13147, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13149 = stablehlo.multiply %13148, %arg496 : tensor<1x8192xf32> loc(#loc5262)
    %13150 = stablehlo.add %13146, %13149 : tensor<1x8192xf32> loc(#loc5263)
    %13151 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13152 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13153 = stablehlo.broadcast_in_dim %13152, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13154 = stablehlo.add %13151, %13153 : tensor<1xf32> loc(#loc5259)
    %13155 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13156 = stablehlo.broadcast_in_dim %13155, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13157 = stablehlo.subtract %13154, %13156 : tensor<1xf32> loc(#loc5260)
    %13158 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13159 = stablehlo.broadcast_in_dim %13158, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13160 = stablehlo.power %13159, %13157 : tensor<1xf32> loc(#loc5261)
    %13161 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13162 = stablehlo.broadcast_in_dim %13161, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13163 = stablehlo.subtract %13162, %13160 : tensor<1xf32> loc(#loc5260)
    %13164 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13165 = stablehlo.broadcast_in_dim %13164, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13166 = stablehlo.multiply %13165, %13163 : tensor<1xf32> loc(#loc5236)
    %13167 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13168 = stablehlo.broadcast_in_dim %13167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13169 = stablehlo.power %13168, %13154 : tensor<1xf32> loc(#loc5261)
    %13170 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13171 = stablehlo.broadcast_in_dim %13170, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13172 = stablehlo.subtract %13171, %13169 : tensor<1xf32> loc(#loc5260)
    %13173 = stablehlo.divide %13166, %13172 : tensor<1xf32> loc(#loc5253)
    %13174 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13175 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13176 = stablehlo.broadcast_in_dim %13175, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13177 = stablehlo.add %13174, %13176 : tensor<1xf32> loc(#loc5259)
    %13178 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13179 = stablehlo.broadcast_in_dim %13178, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13180 = stablehlo.subtract %13177, %13179 : tensor<1xf32> loc(#loc5260)
    %13181 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13182 = stablehlo.broadcast_in_dim %13181, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13183 = stablehlo.power %13182, %13180 : tensor<1xf32> loc(#loc5261)
    %13184 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13185 = stablehlo.broadcast_in_dim %13184, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13186 = stablehlo.subtract %13185, %13183 : tensor<1xf32> loc(#loc5260)
    %13187 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13188 = stablehlo.broadcast_in_dim %13187, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13189 = stablehlo.multiply %13188, %13186 : tensor<1xf32> loc(#loc5236)
    %13190 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13191 = stablehlo.broadcast_in_dim %13190, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13192 = stablehlo.power %13191, %13177 : tensor<1xf32> loc(#loc5261)
    %13193 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13194 = stablehlo.broadcast_in_dim %13193, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13195 = stablehlo.subtract %13194, %13192 : tensor<1xf32> loc(#loc5260)
    %13196 = stablehlo.divide %13189, %13195 : tensor<1xf32> loc(#loc5253)
    %13197 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13198 = stablehlo.broadcast_in_dim %13197, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13199 = stablehlo.subtract %13198, %13173 : tensor<1xf32> loc(#loc5260)
    %13200 = stablehlo.broadcast_in_dim %13199, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13201 = stablehlo.broadcast_in_dim %13200, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13202 = stablehlo.multiply %13201, %4764 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13203 = stablehlo.broadcast_in_dim %13173, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13204 = stablehlo.broadcast_in_dim %13203, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13205 = stablehlo.multiply %13204, %arg321 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13206 = stablehlo.add %13202, %13205 : tensor<1x2048x8192xf32> loc(#loc5263)
    %13207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13208 = stablehlo.broadcast_in_dim %13207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13209 = stablehlo.subtract %13208, %13196 : tensor<1xf32> loc(#loc5260)
    %13210 = stablehlo.multiply %4764, %4764 : tensor<1x2048x8192xf32> loc(#loc5264)
    %13211 = stablehlo.broadcast_in_dim %13209, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13212 = stablehlo.broadcast_in_dim %13211, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13213 = stablehlo.multiply %13212, %13210 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13214 = stablehlo.broadcast_in_dim %13196, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13215 = stablehlo.broadcast_in_dim %13214, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13216 = stablehlo.multiply %13215, %arg497 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13217 = stablehlo.add %13213, %13216 : tensor<1x2048x8192xf32> loc(#loc5263)
    %13218 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13219 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13220 = stablehlo.broadcast_in_dim %13219, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13221 = stablehlo.add %13218, %13220 : tensor<1xf32> loc(#loc5259)
    %13222 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13223 = stablehlo.broadcast_in_dim %13222, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13224 = stablehlo.subtract %13221, %13223 : tensor<1xf32> loc(#loc5260)
    %13225 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13226 = stablehlo.broadcast_in_dim %13225, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13227 = stablehlo.power %13226, %13224 : tensor<1xf32> loc(#loc5261)
    %13228 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13229 = stablehlo.broadcast_in_dim %13228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13230 = stablehlo.subtract %13229, %13227 : tensor<1xf32> loc(#loc5260)
    %13231 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13232 = stablehlo.broadcast_in_dim %13231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13233 = stablehlo.multiply %13232, %13230 : tensor<1xf32> loc(#loc5236)
    %13234 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13235 = stablehlo.broadcast_in_dim %13234, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13236 = stablehlo.power %13235, %13221 : tensor<1xf32> loc(#loc5261)
    %13237 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13238 = stablehlo.broadcast_in_dim %13237, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13239 = stablehlo.subtract %13238, %13236 : tensor<1xf32> loc(#loc5260)
    %13240 = stablehlo.divide %13233, %13239 : tensor<1xf32> loc(#loc5253)
    %13241 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13242 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13243 = stablehlo.broadcast_in_dim %13242, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13244 = stablehlo.add %13241, %13243 : tensor<1xf32> loc(#loc5259)
    %13245 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13246 = stablehlo.broadcast_in_dim %13245, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13247 = stablehlo.subtract %13244, %13246 : tensor<1xf32> loc(#loc5260)
    %13248 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13249 = stablehlo.broadcast_in_dim %13248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13250 = stablehlo.power %13249, %13247 : tensor<1xf32> loc(#loc5261)
    %13251 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13252 = stablehlo.broadcast_in_dim %13251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13253 = stablehlo.subtract %13252, %13250 : tensor<1xf32> loc(#loc5260)
    %13254 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13255 = stablehlo.broadcast_in_dim %13254, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13256 = stablehlo.multiply %13255, %13253 : tensor<1xf32> loc(#loc5236)
    %13257 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13258 = stablehlo.broadcast_in_dim %13257, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13259 = stablehlo.power %13258, %13244 : tensor<1xf32> loc(#loc5261)
    %13260 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13261 = stablehlo.broadcast_in_dim %13260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13262 = stablehlo.subtract %13261, %13259 : tensor<1xf32> loc(#loc5260)
    %13263 = stablehlo.divide %13256, %13262 : tensor<1xf32> loc(#loc5253)
    %13264 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13265 = stablehlo.broadcast_in_dim %13264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13266 = stablehlo.subtract %13265, %13240 : tensor<1xf32> loc(#loc5260)
    %13267 = stablehlo.broadcast_in_dim %13266, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13268 = stablehlo.broadcast_in_dim %13267, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13269 = stablehlo.multiply %13268, %4768 : tensor<1x2048xf32> loc(#loc5262)
    %13270 = stablehlo.broadcast_in_dim %13240, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13271 = stablehlo.broadcast_in_dim %13270, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13272 = stablehlo.multiply %13271, %arg322 : tensor<1x2048xf32> loc(#loc5262)
    %13273 = stablehlo.add %13269, %13272 : tensor<1x2048xf32> loc(#loc5263)
    %13274 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13275 = stablehlo.broadcast_in_dim %13274, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13276 = stablehlo.subtract %13275, %13263 : tensor<1xf32> loc(#loc5260)
    %13277 = stablehlo.multiply %4768, %4768 : tensor<1x2048xf32> loc(#loc5264)
    %13278 = stablehlo.broadcast_in_dim %13276, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13279 = stablehlo.broadcast_in_dim %13278, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13280 = stablehlo.multiply %13279, %13277 : tensor<1x2048xf32> loc(#loc5262)
    %13281 = stablehlo.broadcast_in_dim %13263, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13282 = stablehlo.broadcast_in_dim %13281, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13283 = stablehlo.multiply %13282, %arg498 : tensor<1x2048xf32> loc(#loc5262)
    %13284 = stablehlo.add %13280, %13283 : tensor<1x2048xf32> loc(#loc5263)
    %13285 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13286 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13287 = stablehlo.broadcast_in_dim %13286, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13288 = stablehlo.add %13285, %13287 : tensor<1xf32> loc(#loc5259)
    %13289 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13290 = stablehlo.broadcast_in_dim %13289, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13291 = stablehlo.subtract %13288, %13290 : tensor<1xf32> loc(#loc5260)
    %13292 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13293 = stablehlo.broadcast_in_dim %13292, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13294 = stablehlo.power %13293, %13291 : tensor<1xf32> loc(#loc5261)
    %13295 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13296 = stablehlo.broadcast_in_dim %13295, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13297 = stablehlo.subtract %13296, %13294 : tensor<1xf32> loc(#loc5260)
    %13298 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13299 = stablehlo.broadcast_in_dim %13298, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13300 = stablehlo.multiply %13299, %13297 : tensor<1xf32> loc(#loc5236)
    %13301 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13302 = stablehlo.broadcast_in_dim %13301, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13303 = stablehlo.power %13302, %13288 : tensor<1xf32> loc(#loc5261)
    %13304 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13305 = stablehlo.broadcast_in_dim %13304, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13306 = stablehlo.subtract %13305, %13303 : tensor<1xf32> loc(#loc5260)
    %13307 = stablehlo.divide %13300, %13306 : tensor<1xf32> loc(#loc5253)
    %13308 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13309 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13310 = stablehlo.broadcast_in_dim %13309, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13311 = stablehlo.add %13308, %13310 : tensor<1xf32> loc(#loc5259)
    %13312 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13313 = stablehlo.broadcast_in_dim %13312, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13314 = stablehlo.subtract %13311, %13313 : tensor<1xf32> loc(#loc5260)
    %13315 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13316 = stablehlo.broadcast_in_dim %13315, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13317 = stablehlo.power %13316, %13314 : tensor<1xf32> loc(#loc5261)
    %13318 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13319 = stablehlo.broadcast_in_dim %13318, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13320 = stablehlo.subtract %13319, %13317 : tensor<1xf32> loc(#loc5260)
    %13321 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13322 = stablehlo.broadcast_in_dim %13321, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13323 = stablehlo.multiply %13322, %13320 : tensor<1xf32> loc(#loc5236)
    %13324 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13325 = stablehlo.broadcast_in_dim %13324, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13326 = stablehlo.power %13325, %13311 : tensor<1xf32> loc(#loc5261)
    %13327 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13328 = stablehlo.broadcast_in_dim %13327, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13329 = stablehlo.subtract %13328, %13326 : tensor<1xf32> loc(#loc5260)
    %13330 = stablehlo.divide %13323, %13329 : tensor<1xf32> loc(#loc5253)
    %13331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13332 = stablehlo.broadcast_in_dim %13331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13333 = stablehlo.subtract %13332, %13307 : tensor<1xf32> loc(#loc5260)
    %13334 = stablehlo.broadcast_in_dim %13333, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13335 = stablehlo.broadcast_in_dim %13334, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %13336 = stablehlo.multiply %13335, %4772 : tensor<1x8192x2048xf32> loc(#loc5262)
    %13337 = stablehlo.broadcast_in_dim %13307, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13338 = stablehlo.broadcast_in_dim %13337, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %13339 = stablehlo.multiply %13338, %arg323 : tensor<1x8192x2048xf32> loc(#loc5262)
    %13340 = stablehlo.add %13336, %13339 : tensor<1x8192x2048xf32> loc(#loc5263)
    %13341 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13342 = stablehlo.broadcast_in_dim %13341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13343 = stablehlo.subtract %13342, %13330 : tensor<1xf32> loc(#loc5260)
    %13344 = stablehlo.multiply %4772, %4772 : tensor<1x8192x2048xf32> loc(#loc5264)
    %13345 = stablehlo.broadcast_in_dim %13343, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13346 = stablehlo.broadcast_in_dim %13345, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %13347 = stablehlo.multiply %13346, %13344 : tensor<1x8192x2048xf32> loc(#loc5262)
    %13348 = stablehlo.broadcast_in_dim %13330, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13349 = stablehlo.broadcast_in_dim %13348, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %13350 = stablehlo.multiply %13349, %arg499 : tensor<1x8192x2048xf32> loc(#loc5262)
    %13351 = stablehlo.add %13347, %13350 : tensor<1x8192x2048xf32> loc(#loc5263)
    %13352 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13353 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13354 = stablehlo.broadcast_in_dim %13353, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13355 = stablehlo.add %13352, %13354 : tensor<1xf32> loc(#loc5259)
    %13356 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13357 = stablehlo.broadcast_in_dim %13356, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13358 = stablehlo.subtract %13355, %13357 : tensor<1xf32> loc(#loc5260)
    %13359 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13360 = stablehlo.broadcast_in_dim %13359, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13361 = stablehlo.power %13360, %13358 : tensor<1xf32> loc(#loc5261)
    %13362 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13363 = stablehlo.broadcast_in_dim %13362, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13364 = stablehlo.subtract %13363, %13361 : tensor<1xf32> loc(#loc5260)
    %13365 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13366 = stablehlo.broadcast_in_dim %13365, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13367 = stablehlo.multiply %13366, %13364 : tensor<1xf32> loc(#loc5236)
    %13368 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13369 = stablehlo.broadcast_in_dim %13368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13370 = stablehlo.power %13369, %13355 : tensor<1xf32> loc(#loc5261)
    %13371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13372 = stablehlo.broadcast_in_dim %13371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13373 = stablehlo.subtract %13372, %13370 : tensor<1xf32> loc(#loc5260)
    %13374 = stablehlo.divide %13367, %13373 : tensor<1xf32> loc(#loc5253)
    %13375 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13376 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13377 = stablehlo.broadcast_in_dim %13376, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13378 = stablehlo.add %13375, %13377 : tensor<1xf32> loc(#loc5259)
    %13379 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13380 = stablehlo.broadcast_in_dim %13379, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13381 = stablehlo.subtract %13378, %13380 : tensor<1xf32> loc(#loc5260)
    %13382 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13383 = stablehlo.broadcast_in_dim %13382, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13384 = stablehlo.power %13383, %13381 : tensor<1xf32> loc(#loc5261)
    %13385 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13386 = stablehlo.broadcast_in_dim %13385, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13387 = stablehlo.subtract %13386, %13384 : tensor<1xf32> loc(#loc5260)
    %13388 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13389 = stablehlo.broadcast_in_dim %13388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13390 = stablehlo.multiply %13389, %13387 : tensor<1xf32> loc(#loc5236)
    %13391 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13392 = stablehlo.broadcast_in_dim %13391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13393 = stablehlo.power %13392, %13378 : tensor<1xf32> loc(#loc5261)
    %13394 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13395 = stablehlo.broadcast_in_dim %13394, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13396 = stablehlo.subtract %13395, %13393 : tensor<1xf32> loc(#loc5260)
    %13397 = stablehlo.divide %13390, %13396 : tensor<1xf32> loc(#loc5253)
    %13398 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13399 = stablehlo.broadcast_in_dim %13398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13400 = stablehlo.subtract %13399, %13374 : tensor<1xf32> loc(#loc5260)
    %13401 = stablehlo.broadcast_in_dim %13400, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13402 = stablehlo.broadcast_in_dim %13401, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13403 = stablehlo.multiply %13402, %4776 : tensor<1x2048xf32> loc(#loc5262)
    %13404 = stablehlo.broadcast_in_dim %13374, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13405 = stablehlo.broadcast_in_dim %13404, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13406 = stablehlo.multiply %13405, %arg324 : tensor<1x2048xf32> loc(#loc5262)
    %13407 = stablehlo.add %13403, %13406 : tensor<1x2048xf32> loc(#loc5263)
    %13408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13409 = stablehlo.broadcast_in_dim %13408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13410 = stablehlo.subtract %13409, %13397 : tensor<1xf32> loc(#loc5260)
    %13411 = stablehlo.multiply %4776, %4776 : tensor<1x2048xf32> loc(#loc5264)
    %13412 = stablehlo.broadcast_in_dim %13410, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13413 = stablehlo.broadcast_in_dim %13412, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13414 = stablehlo.multiply %13413, %13411 : tensor<1x2048xf32> loc(#loc5262)
    %13415 = stablehlo.broadcast_in_dim %13397, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13416 = stablehlo.broadcast_in_dim %13415, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13417 = stablehlo.multiply %13416, %arg500 : tensor<1x2048xf32> loc(#loc5262)
    %13418 = stablehlo.add %13414, %13417 : tensor<1x2048xf32> loc(#loc5263)
    %13419 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13420 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13421 = stablehlo.broadcast_in_dim %13420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13422 = stablehlo.add %13419, %13421 : tensor<1xf32> loc(#loc5259)
    %13423 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13424 = stablehlo.broadcast_in_dim %13423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13425 = stablehlo.subtract %13422, %13424 : tensor<1xf32> loc(#loc5260)
    %13426 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13427 = stablehlo.broadcast_in_dim %13426, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13428 = stablehlo.power %13427, %13425 : tensor<1xf32> loc(#loc5261)
    %13429 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13430 = stablehlo.broadcast_in_dim %13429, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13431 = stablehlo.subtract %13430, %13428 : tensor<1xf32> loc(#loc5260)
    %13432 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13433 = stablehlo.broadcast_in_dim %13432, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13434 = stablehlo.multiply %13433, %13431 : tensor<1xf32> loc(#loc5236)
    %13435 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13436 = stablehlo.broadcast_in_dim %13435, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13437 = stablehlo.power %13436, %13422 : tensor<1xf32> loc(#loc5261)
    %13438 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13439 = stablehlo.broadcast_in_dim %13438, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13440 = stablehlo.subtract %13439, %13437 : tensor<1xf32> loc(#loc5260)
    %13441 = stablehlo.divide %13434, %13440 : tensor<1xf32> loc(#loc5253)
    %13442 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13443 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13444 = stablehlo.broadcast_in_dim %13443, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13445 = stablehlo.add %13442, %13444 : tensor<1xf32> loc(#loc5259)
    %13446 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13447 = stablehlo.broadcast_in_dim %13446, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13448 = stablehlo.subtract %13445, %13447 : tensor<1xf32> loc(#loc5260)
    %13449 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13450 = stablehlo.broadcast_in_dim %13449, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13451 = stablehlo.power %13450, %13448 : tensor<1xf32> loc(#loc5261)
    %13452 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13453 = stablehlo.broadcast_in_dim %13452, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13454 = stablehlo.subtract %13453, %13451 : tensor<1xf32> loc(#loc5260)
    %13455 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13456 = stablehlo.broadcast_in_dim %13455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13457 = stablehlo.multiply %13456, %13454 : tensor<1xf32> loc(#loc5236)
    %13458 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13459 = stablehlo.broadcast_in_dim %13458, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13460 = stablehlo.power %13459, %13445 : tensor<1xf32> loc(#loc5261)
    %13461 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13462 = stablehlo.broadcast_in_dim %13461, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13463 = stablehlo.subtract %13462, %13460 : tensor<1xf32> loc(#loc5260)
    %13464 = stablehlo.divide %13457, %13463 : tensor<1xf32> loc(#loc5253)
    %13465 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13466 = stablehlo.broadcast_in_dim %13465, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13467 = stablehlo.subtract %13466, %13441 : tensor<1xf32> loc(#loc5260)
    %13468 = stablehlo.broadcast_in_dim %13467, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13469 = stablehlo.broadcast_in_dim %13468, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13470 = stablehlo.multiply %13469, %4780 : tensor<1x2048xf32> loc(#loc5262)
    %13471 = stablehlo.broadcast_in_dim %13441, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13472 = stablehlo.broadcast_in_dim %13471, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13473 = stablehlo.multiply %13472, %arg325 : tensor<1x2048xf32> loc(#loc5262)
    %13474 = stablehlo.add %13470, %13473 : tensor<1x2048xf32> loc(#loc5263)
    %13475 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13476 = stablehlo.broadcast_in_dim %13475, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13477 = stablehlo.subtract %13476, %13464 : tensor<1xf32> loc(#loc5260)
    %13478 = stablehlo.multiply %4780, %4780 : tensor<1x2048xf32> loc(#loc5264)
    %13479 = stablehlo.broadcast_in_dim %13477, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13480 = stablehlo.broadcast_in_dim %13479, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13481 = stablehlo.multiply %13480, %13478 : tensor<1x2048xf32> loc(#loc5262)
    %13482 = stablehlo.broadcast_in_dim %13464, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13483 = stablehlo.broadcast_in_dim %13482, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13484 = stablehlo.multiply %13483, %arg501 : tensor<1x2048xf32> loc(#loc5262)
    %13485 = stablehlo.add %13481, %13484 : tensor<1x2048xf32> loc(#loc5263)
    %13486 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13487 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13488 = stablehlo.broadcast_in_dim %13487, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13489 = stablehlo.add %13486, %13488 : tensor<1xf32> loc(#loc5259)
    %13490 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13491 = stablehlo.broadcast_in_dim %13490, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13492 = stablehlo.subtract %13489, %13491 : tensor<1xf32> loc(#loc5260)
    %13493 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13494 = stablehlo.broadcast_in_dim %13493, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13495 = stablehlo.power %13494, %13492 : tensor<1xf32> loc(#loc5261)
    %13496 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13497 = stablehlo.broadcast_in_dim %13496, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13498 = stablehlo.subtract %13497, %13495 : tensor<1xf32> loc(#loc5260)
    %13499 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13500 = stablehlo.broadcast_in_dim %13499, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13501 = stablehlo.multiply %13500, %13498 : tensor<1xf32> loc(#loc5236)
    %13502 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13503 = stablehlo.broadcast_in_dim %13502, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13504 = stablehlo.power %13503, %13489 : tensor<1xf32> loc(#loc5261)
    %13505 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13506 = stablehlo.broadcast_in_dim %13505, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13507 = stablehlo.subtract %13506, %13504 : tensor<1xf32> loc(#loc5260)
    %13508 = stablehlo.divide %13501, %13507 : tensor<1xf32> loc(#loc5253)
    %13509 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13510 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13511 = stablehlo.broadcast_in_dim %13510, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13512 = stablehlo.add %13509, %13511 : tensor<1xf32> loc(#loc5259)
    %13513 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13514 = stablehlo.broadcast_in_dim %13513, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13515 = stablehlo.subtract %13512, %13514 : tensor<1xf32> loc(#loc5260)
    %13516 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13517 = stablehlo.broadcast_in_dim %13516, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13518 = stablehlo.power %13517, %13515 : tensor<1xf32> loc(#loc5261)
    %13519 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13520 = stablehlo.broadcast_in_dim %13519, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13521 = stablehlo.subtract %13520, %13518 : tensor<1xf32> loc(#loc5260)
    %13522 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13523 = stablehlo.broadcast_in_dim %13522, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13524 = stablehlo.multiply %13523, %13521 : tensor<1xf32> loc(#loc5236)
    %13525 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13526 = stablehlo.broadcast_in_dim %13525, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13527 = stablehlo.power %13526, %13512 : tensor<1xf32> loc(#loc5261)
    %13528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13529 = stablehlo.broadcast_in_dim %13528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13530 = stablehlo.subtract %13529, %13527 : tensor<1xf32> loc(#loc5260)
    %13531 = stablehlo.divide %13524, %13530 : tensor<1xf32> loc(#loc5253)
    %13532 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13533 = stablehlo.broadcast_in_dim %13532, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13534 = stablehlo.subtract %13533, %13508 : tensor<1xf32> loc(#loc5260)
    %13535 = stablehlo.broadcast_in_dim %13534, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13536 = stablehlo.broadcast_in_dim %13535, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13537 = stablehlo.multiply %13536, %4784 : tensor<1x2048xf32> loc(#loc5262)
    %13538 = stablehlo.broadcast_in_dim %13508, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13539 = stablehlo.broadcast_in_dim %13538, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13540 = stablehlo.multiply %13539, %arg326 : tensor<1x2048xf32> loc(#loc5262)
    %13541 = stablehlo.add %13537, %13540 : tensor<1x2048xf32> loc(#loc5263)
    %13542 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13543 = stablehlo.broadcast_in_dim %13542, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13544 = stablehlo.subtract %13543, %13531 : tensor<1xf32> loc(#loc5260)
    %13545 = stablehlo.multiply %4784, %4784 : tensor<1x2048xf32> loc(#loc5264)
    %13546 = stablehlo.broadcast_in_dim %13544, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13547 = stablehlo.broadcast_in_dim %13546, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13548 = stablehlo.multiply %13547, %13545 : tensor<1x2048xf32> loc(#loc5262)
    %13549 = stablehlo.broadcast_in_dim %13531, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13550 = stablehlo.broadcast_in_dim %13549, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13551 = stablehlo.multiply %13550, %arg502 : tensor<1x2048xf32> loc(#loc5262)
    %13552 = stablehlo.add %13548, %13551 : tensor<1x2048xf32> loc(#loc5263)
    %13553 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13554 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13555 = stablehlo.broadcast_in_dim %13554, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13556 = stablehlo.add %13553, %13555 : tensor<1xf32> loc(#loc5259)
    %13557 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13558 = stablehlo.broadcast_in_dim %13557, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13559 = stablehlo.subtract %13556, %13558 : tensor<1xf32> loc(#loc5260)
    %13560 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13561 = stablehlo.broadcast_in_dim %13560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13562 = stablehlo.power %13561, %13559 : tensor<1xf32> loc(#loc5261)
    %13563 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13564 = stablehlo.broadcast_in_dim %13563, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13565 = stablehlo.subtract %13564, %13562 : tensor<1xf32> loc(#loc5260)
    %13566 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13567 = stablehlo.broadcast_in_dim %13566, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13568 = stablehlo.multiply %13567, %13565 : tensor<1xf32> loc(#loc5236)
    %13569 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13570 = stablehlo.broadcast_in_dim %13569, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13571 = stablehlo.power %13570, %13556 : tensor<1xf32> loc(#loc5261)
    %13572 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13573 = stablehlo.broadcast_in_dim %13572, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13574 = stablehlo.subtract %13573, %13571 : tensor<1xf32> loc(#loc5260)
    %13575 = stablehlo.divide %13568, %13574 : tensor<1xf32> loc(#loc5253)
    %13576 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13577 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13578 = stablehlo.broadcast_in_dim %13577, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13579 = stablehlo.add %13576, %13578 : tensor<1xf32> loc(#loc5259)
    %13580 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13581 = stablehlo.broadcast_in_dim %13580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13582 = stablehlo.subtract %13579, %13581 : tensor<1xf32> loc(#loc5260)
    %13583 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13584 = stablehlo.broadcast_in_dim %13583, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13585 = stablehlo.power %13584, %13582 : tensor<1xf32> loc(#loc5261)
    %13586 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13587 = stablehlo.broadcast_in_dim %13586, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13588 = stablehlo.subtract %13587, %13585 : tensor<1xf32> loc(#loc5260)
    %13589 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13590 = stablehlo.broadcast_in_dim %13589, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13591 = stablehlo.multiply %13590, %13588 : tensor<1xf32> loc(#loc5236)
    %13592 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13593 = stablehlo.broadcast_in_dim %13592, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13594 = stablehlo.power %13593, %13579 : tensor<1xf32> loc(#loc5261)
    %13595 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13596 = stablehlo.broadcast_in_dim %13595, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13597 = stablehlo.subtract %13596, %13594 : tensor<1xf32> loc(#loc5260)
    %13598 = stablehlo.divide %13591, %13597 : tensor<1xf32> loc(#loc5253)
    %13599 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13600 = stablehlo.broadcast_in_dim %13599, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13601 = stablehlo.subtract %13600, %13575 : tensor<1xf32> loc(#loc5260)
    %13602 = stablehlo.broadcast_in_dim %13601, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13603 = stablehlo.broadcast_in_dim %13602, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13604 = stablehlo.multiply %13603, %4788 : tensor<1x2048xf32> loc(#loc5262)
    %13605 = stablehlo.broadcast_in_dim %13575, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13606 = stablehlo.broadcast_in_dim %13605, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13607 = stablehlo.multiply %13606, %arg327 : tensor<1x2048xf32> loc(#loc5262)
    %13608 = stablehlo.add %13604, %13607 : tensor<1x2048xf32> loc(#loc5263)
    %13609 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13610 = stablehlo.broadcast_in_dim %13609, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13611 = stablehlo.subtract %13610, %13598 : tensor<1xf32> loc(#loc5260)
    %13612 = stablehlo.multiply %4788, %4788 : tensor<1x2048xf32> loc(#loc5264)
    %13613 = stablehlo.broadcast_in_dim %13611, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13614 = stablehlo.broadcast_in_dim %13613, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13615 = stablehlo.multiply %13614, %13612 : tensor<1x2048xf32> loc(#loc5262)
    %13616 = stablehlo.broadcast_in_dim %13598, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13617 = stablehlo.broadcast_in_dim %13616, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %13618 = stablehlo.multiply %13617, %arg503 : tensor<1x2048xf32> loc(#loc5262)
    %13619 = stablehlo.add %13615, %13618 : tensor<1x2048xf32> loc(#loc5263)
    %13620 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13621 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13622 = stablehlo.broadcast_in_dim %13621, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13623 = stablehlo.add %13620, %13622 : tensor<1xf32> loc(#loc5259)
    %13624 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13625 = stablehlo.broadcast_in_dim %13624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13626 = stablehlo.subtract %13623, %13625 : tensor<1xf32> loc(#loc5260)
    %13627 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13628 = stablehlo.broadcast_in_dim %13627, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13629 = stablehlo.power %13628, %13626 : tensor<1xf32> loc(#loc5261)
    %13630 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13631 = stablehlo.broadcast_in_dim %13630, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13632 = stablehlo.subtract %13631, %13629 : tensor<1xf32> loc(#loc5260)
    %13633 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13634 = stablehlo.broadcast_in_dim %13633, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13635 = stablehlo.multiply %13634, %13632 : tensor<1xf32> loc(#loc5236)
    %13636 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13637 = stablehlo.broadcast_in_dim %13636, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13638 = stablehlo.power %13637, %13623 : tensor<1xf32> loc(#loc5261)
    %13639 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13640 = stablehlo.broadcast_in_dim %13639, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13641 = stablehlo.subtract %13640, %13638 : tensor<1xf32> loc(#loc5260)
    %13642 = stablehlo.divide %13635, %13641 : tensor<1xf32> loc(#loc5253)
    %13643 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13644 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13645 = stablehlo.broadcast_in_dim %13644, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13646 = stablehlo.add %13643, %13645 : tensor<1xf32> loc(#loc5259)
    %13647 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13648 = stablehlo.broadcast_in_dim %13647, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13649 = stablehlo.subtract %13646, %13648 : tensor<1xf32> loc(#loc5260)
    %13650 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13651 = stablehlo.broadcast_in_dim %13650, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13652 = stablehlo.power %13651, %13649 : tensor<1xf32> loc(#loc5261)
    %13653 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13654 = stablehlo.broadcast_in_dim %13653, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13655 = stablehlo.subtract %13654, %13652 : tensor<1xf32> loc(#loc5260)
    %13656 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13657 = stablehlo.broadcast_in_dim %13656, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13658 = stablehlo.multiply %13657, %13655 : tensor<1xf32> loc(#loc5236)
    %13659 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13660 = stablehlo.broadcast_in_dim %13659, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13661 = stablehlo.power %13660, %13646 : tensor<1xf32> loc(#loc5261)
    %13662 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13663 = stablehlo.broadcast_in_dim %13662, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13664 = stablehlo.subtract %13663, %13661 : tensor<1xf32> loc(#loc5260)
    %13665 = stablehlo.divide %13658, %13664 : tensor<1xf32> loc(#loc5253)
    %13666 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13667 = stablehlo.broadcast_in_dim %13666, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13668 = stablehlo.subtract %13667, %13642 : tensor<1xf32> loc(#loc5260)
    %13669 = stablehlo.broadcast_in_dim %13668, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %13670 = stablehlo.broadcast_in_dim %13669, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13671 = stablehlo.multiply %13670, %4792 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13672 = stablehlo.broadcast_in_dim %13642, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %13673 = stablehlo.broadcast_in_dim %13672, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13674 = stablehlo.multiply %13673, %arg328 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13675 = stablehlo.add %13671, %13674 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %13676 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13677 = stablehlo.broadcast_in_dim %13676, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13678 = stablehlo.subtract %13677, %13665 : tensor<1xf32> loc(#loc5260)
    %13679 = stablehlo.multiply %4792, %4792 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %13680 = stablehlo.broadcast_in_dim %13678, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %13681 = stablehlo.broadcast_in_dim %13680, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13682 = stablehlo.multiply %13681, %13679 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13683 = stablehlo.broadcast_in_dim %13665, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %13684 = stablehlo.broadcast_in_dim %13683, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13685 = stablehlo.multiply %13684, %arg504 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %13686 = stablehlo.add %13682, %13685 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %13687 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13688 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13689 = stablehlo.broadcast_in_dim %13688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13690 = stablehlo.add %13687, %13689 : tensor<1xf32> loc(#loc5259)
    %13691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13692 = stablehlo.broadcast_in_dim %13691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13693 = stablehlo.subtract %13690, %13692 : tensor<1xf32> loc(#loc5260)
    %13694 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13695 = stablehlo.broadcast_in_dim %13694, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13696 = stablehlo.power %13695, %13693 : tensor<1xf32> loc(#loc5261)
    %13697 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13698 = stablehlo.broadcast_in_dim %13697, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13699 = stablehlo.subtract %13698, %13696 : tensor<1xf32> loc(#loc5260)
    %13700 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13701 = stablehlo.broadcast_in_dim %13700, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13702 = stablehlo.multiply %13701, %13699 : tensor<1xf32> loc(#loc5236)
    %13703 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13704 = stablehlo.broadcast_in_dim %13703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13705 = stablehlo.power %13704, %13690 : tensor<1xf32> loc(#loc5261)
    %13706 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13707 = stablehlo.broadcast_in_dim %13706, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13708 = stablehlo.subtract %13707, %13705 : tensor<1xf32> loc(#loc5260)
    %13709 = stablehlo.divide %13702, %13708 : tensor<1xf32> loc(#loc5253)
    %13710 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13712 = stablehlo.broadcast_in_dim %13711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13713 = stablehlo.add %13710, %13712 : tensor<1xf32> loc(#loc5259)
    %13714 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13715 = stablehlo.broadcast_in_dim %13714, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13716 = stablehlo.subtract %13713, %13715 : tensor<1xf32> loc(#loc5260)
    %13717 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13718 = stablehlo.broadcast_in_dim %13717, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13719 = stablehlo.power %13718, %13716 : tensor<1xf32> loc(#loc5261)
    %13720 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13721 = stablehlo.broadcast_in_dim %13720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13722 = stablehlo.subtract %13721, %13719 : tensor<1xf32> loc(#loc5260)
    %13723 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13724 = stablehlo.broadcast_in_dim %13723, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13725 = stablehlo.multiply %13724, %13722 : tensor<1xf32> loc(#loc5236)
    %13726 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13727 = stablehlo.broadcast_in_dim %13726, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13728 = stablehlo.power %13727, %13713 : tensor<1xf32> loc(#loc5261)
    %13729 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13730 = stablehlo.broadcast_in_dim %13729, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13731 = stablehlo.subtract %13730, %13728 : tensor<1xf32> loc(#loc5260)
    %13732 = stablehlo.divide %13725, %13731 : tensor<1xf32> loc(#loc5253)
    %13733 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13734 = stablehlo.broadcast_in_dim %13733, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13735 = stablehlo.subtract %13734, %13709 : tensor<1xf32> loc(#loc5260)
    %13736 = stablehlo.broadcast_in_dim %13735, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13737 = stablehlo.broadcast_in_dim %13736, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13738 = stablehlo.multiply %13737, %4796 : tensor<1x64xf32> loc(#loc5262)
    %13739 = stablehlo.broadcast_in_dim %13709, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13740 = stablehlo.broadcast_in_dim %13739, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13741 = stablehlo.multiply %13740, %arg329 : tensor<1x64xf32> loc(#loc5262)
    %13742 = stablehlo.add %13738, %13741 : tensor<1x64xf32> loc(#loc5263)
    %13743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13744 = stablehlo.broadcast_in_dim %13743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13745 = stablehlo.subtract %13744, %13732 : tensor<1xf32> loc(#loc5260)
    %13746 = stablehlo.multiply %4796, %4796 : tensor<1x64xf32> loc(#loc5264)
    %13747 = stablehlo.broadcast_in_dim %13745, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13748 = stablehlo.broadcast_in_dim %13747, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13749 = stablehlo.multiply %13748, %13746 : tensor<1x64xf32> loc(#loc5262)
    %13750 = stablehlo.broadcast_in_dim %13732, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13751 = stablehlo.broadcast_in_dim %13750, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %13752 = stablehlo.multiply %13751, %arg505 : tensor<1x64xf32> loc(#loc5262)
    %13753 = stablehlo.add %13749, %13752 : tensor<1x64xf32> loc(#loc5263)
    %13754 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13755 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13756 = stablehlo.broadcast_in_dim %13755, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13757 = stablehlo.add %13754, %13756 : tensor<1xf32> loc(#loc5259)
    %13758 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13759 = stablehlo.broadcast_in_dim %13758, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13760 = stablehlo.subtract %13757, %13759 : tensor<1xf32> loc(#loc5260)
    %13761 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13762 = stablehlo.broadcast_in_dim %13761, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13763 = stablehlo.power %13762, %13760 : tensor<1xf32> loc(#loc5261)
    %13764 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13765 = stablehlo.broadcast_in_dim %13764, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13766 = stablehlo.subtract %13765, %13763 : tensor<1xf32> loc(#loc5260)
    %13767 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13768 = stablehlo.broadcast_in_dim %13767, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13769 = stablehlo.multiply %13768, %13766 : tensor<1xf32> loc(#loc5236)
    %13770 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13771 = stablehlo.broadcast_in_dim %13770, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13772 = stablehlo.power %13771, %13757 : tensor<1xf32> loc(#loc5261)
    %13773 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13774 = stablehlo.broadcast_in_dim %13773, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13775 = stablehlo.subtract %13774, %13772 : tensor<1xf32> loc(#loc5260)
    %13776 = stablehlo.divide %13769, %13775 : tensor<1xf32> loc(#loc5253)
    %13777 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13778 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13779 = stablehlo.broadcast_in_dim %13778, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13780 = stablehlo.add %13777, %13779 : tensor<1xf32> loc(#loc5259)
    %13781 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13782 = stablehlo.broadcast_in_dim %13781, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13783 = stablehlo.subtract %13780, %13782 : tensor<1xf32> loc(#loc5260)
    %13784 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13785 = stablehlo.broadcast_in_dim %13784, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13786 = stablehlo.power %13785, %13783 : tensor<1xf32> loc(#loc5261)
    %13787 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13788 = stablehlo.broadcast_in_dim %13787, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13789 = stablehlo.subtract %13788, %13786 : tensor<1xf32> loc(#loc5260)
    %13790 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13791 = stablehlo.broadcast_in_dim %13790, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13792 = stablehlo.multiply %13791, %13789 : tensor<1xf32> loc(#loc5236)
    %13793 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13794 = stablehlo.broadcast_in_dim %13793, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13795 = stablehlo.power %13794, %13780 : tensor<1xf32> loc(#loc5261)
    %13796 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13797 = stablehlo.broadcast_in_dim %13796, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13798 = stablehlo.subtract %13797, %13795 : tensor<1xf32> loc(#loc5260)
    %13799 = stablehlo.divide %13792, %13798 : tensor<1xf32> loc(#loc5253)
    %13800 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13801 = stablehlo.broadcast_in_dim %13800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13802 = stablehlo.subtract %13801, %13776 : tensor<1xf32> loc(#loc5260)
    %13803 = stablehlo.broadcast_in_dim %13802, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13804 = stablehlo.broadcast_in_dim %13803, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13805 = stablehlo.multiply %13804, %4800 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13806 = stablehlo.broadcast_in_dim %13776, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13807 = stablehlo.broadcast_in_dim %13806, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13808 = stablehlo.multiply %13807, %arg330 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13809 = stablehlo.add %13805, %13808 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %13810 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13811 = stablehlo.broadcast_in_dim %13810, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13812 = stablehlo.subtract %13811, %13799 : tensor<1xf32> loc(#loc5260)
    %13813 = stablehlo.multiply %4800, %4800 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %13814 = stablehlo.broadcast_in_dim %13812, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13815 = stablehlo.broadcast_in_dim %13814, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13816 = stablehlo.multiply %13815, %13813 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13817 = stablehlo.broadcast_in_dim %13799, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %13818 = stablehlo.broadcast_in_dim %13817, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %13819 = stablehlo.multiply %13818, %arg506 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %13820 = stablehlo.add %13816, %13819 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %13821 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13822 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13823 = stablehlo.broadcast_in_dim %13822, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13824 = stablehlo.add %13821, %13823 : tensor<1xf32> loc(#loc5259)
    %13825 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13826 = stablehlo.broadcast_in_dim %13825, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13827 = stablehlo.subtract %13824, %13826 : tensor<1xf32> loc(#loc5260)
    %13828 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13829 = stablehlo.broadcast_in_dim %13828, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13830 = stablehlo.power %13829, %13827 : tensor<1xf32> loc(#loc5261)
    %13831 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13832 = stablehlo.broadcast_in_dim %13831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13833 = stablehlo.subtract %13832, %13830 : tensor<1xf32> loc(#loc5260)
    %13834 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13835 = stablehlo.broadcast_in_dim %13834, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13836 = stablehlo.multiply %13835, %13833 : tensor<1xf32> loc(#loc5236)
    %13837 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13838 = stablehlo.broadcast_in_dim %13837, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13839 = stablehlo.power %13838, %13824 : tensor<1xf32> loc(#loc5261)
    %13840 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13841 = stablehlo.broadcast_in_dim %13840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13842 = stablehlo.subtract %13841, %13839 : tensor<1xf32> loc(#loc5260)
    %13843 = stablehlo.divide %13836, %13842 : tensor<1xf32> loc(#loc5253)
    %13844 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13845 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13846 = stablehlo.broadcast_in_dim %13845, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13847 = stablehlo.add %13844, %13846 : tensor<1xf32> loc(#loc5259)
    %13848 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13849 = stablehlo.broadcast_in_dim %13848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13850 = stablehlo.subtract %13847, %13849 : tensor<1xf32> loc(#loc5260)
    %13851 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13852 = stablehlo.broadcast_in_dim %13851, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13853 = stablehlo.power %13852, %13850 : tensor<1xf32> loc(#loc5261)
    %13854 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13855 = stablehlo.broadcast_in_dim %13854, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13856 = stablehlo.subtract %13855, %13853 : tensor<1xf32> loc(#loc5260)
    %13857 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13858 = stablehlo.broadcast_in_dim %13857, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13859 = stablehlo.multiply %13858, %13856 : tensor<1xf32> loc(#loc5236)
    %13860 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13861 = stablehlo.broadcast_in_dim %13860, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13862 = stablehlo.power %13861, %13847 : tensor<1xf32> loc(#loc5261)
    %13863 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13864 = stablehlo.broadcast_in_dim %13863, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13865 = stablehlo.subtract %13864, %13862 : tensor<1xf32> loc(#loc5260)
    %13866 = stablehlo.divide %13859, %13865 : tensor<1xf32> loc(#loc5253)
    %13867 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13868 = stablehlo.broadcast_in_dim %13867, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13869 = stablehlo.subtract %13868, %13843 : tensor<1xf32> loc(#loc5260)
    %13870 = stablehlo.broadcast_in_dim %13869, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13871 = stablehlo.broadcast_in_dim %13870, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13872 = stablehlo.multiply %13871, %4804 : tensor<1x8192xf32> loc(#loc5262)
    %13873 = stablehlo.broadcast_in_dim %13843, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13874 = stablehlo.broadcast_in_dim %13873, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13875 = stablehlo.multiply %13874, %arg331 : tensor<1x8192xf32> loc(#loc5262)
    %13876 = stablehlo.add %13872, %13875 : tensor<1x8192xf32> loc(#loc5263)
    %13877 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13878 = stablehlo.broadcast_in_dim %13877, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13879 = stablehlo.subtract %13878, %13866 : tensor<1xf32> loc(#loc5260)
    %13880 = stablehlo.multiply %4804, %4804 : tensor<1x8192xf32> loc(#loc5264)
    %13881 = stablehlo.broadcast_in_dim %13879, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13882 = stablehlo.broadcast_in_dim %13881, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13883 = stablehlo.multiply %13882, %13880 : tensor<1x8192xf32> loc(#loc5262)
    %13884 = stablehlo.broadcast_in_dim %13866, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %13885 = stablehlo.broadcast_in_dim %13884, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %13886 = stablehlo.multiply %13885, %arg507 : tensor<1x8192xf32> loc(#loc5262)
    %13887 = stablehlo.add %13883, %13886 : tensor<1x8192xf32> loc(#loc5263)
    %13888 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13889 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13890 = stablehlo.broadcast_in_dim %13889, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13891 = stablehlo.add %13888, %13890 : tensor<1xf32> loc(#loc5259)
    %13892 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13893 = stablehlo.broadcast_in_dim %13892, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13894 = stablehlo.subtract %13891, %13893 : tensor<1xf32> loc(#loc5260)
    %13895 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13896 = stablehlo.broadcast_in_dim %13895, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13897 = stablehlo.power %13896, %13894 : tensor<1xf32> loc(#loc5261)
    %13898 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13899 = stablehlo.broadcast_in_dim %13898, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13900 = stablehlo.subtract %13899, %13897 : tensor<1xf32> loc(#loc5260)
    %13901 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13902 = stablehlo.broadcast_in_dim %13901, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13903 = stablehlo.multiply %13902, %13900 : tensor<1xf32> loc(#loc5236)
    %13904 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13905 = stablehlo.broadcast_in_dim %13904, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13906 = stablehlo.power %13905, %13891 : tensor<1xf32> loc(#loc5261)
    %13907 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13908 = stablehlo.broadcast_in_dim %13907, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13909 = stablehlo.subtract %13908, %13906 : tensor<1xf32> loc(#loc5260)
    %13910 = stablehlo.divide %13903, %13909 : tensor<1xf32> loc(#loc5253)
    %13911 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13912 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13913 = stablehlo.broadcast_in_dim %13912, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13914 = stablehlo.add %13911, %13913 : tensor<1xf32> loc(#loc5259)
    %13915 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13916 = stablehlo.broadcast_in_dim %13915, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13917 = stablehlo.subtract %13914, %13916 : tensor<1xf32> loc(#loc5260)
    %13918 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13919 = stablehlo.broadcast_in_dim %13918, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13920 = stablehlo.power %13919, %13917 : tensor<1xf32> loc(#loc5261)
    %13921 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13922 = stablehlo.broadcast_in_dim %13921, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13923 = stablehlo.subtract %13922, %13920 : tensor<1xf32> loc(#loc5260)
    %13924 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13925 = stablehlo.broadcast_in_dim %13924, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13926 = stablehlo.multiply %13925, %13923 : tensor<1xf32> loc(#loc5236)
    %13927 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13928 = stablehlo.broadcast_in_dim %13927, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13929 = stablehlo.power %13928, %13914 : tensor<1xf32> loc(#loc5261)
    %13930 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13931 = stablehlo.broadcast_in_dim %13930, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13932 = stablehlo.subtract %13931, %13929 : tensor<1xf32> loc(#loc5260)
    %13933 = stablehlo.divide %13926, %13932 : tensor<1xf32> loc(#loc5253)
    %13934 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13935 = stablehlo.broadcast_in_dim %13934, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13936 = stablehlo.subtract %13935, %13910 : tensor<1xf32> loc(#loc5260)
    %13937 = stablehlo.broadcast_in_dim %13936, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13938 = stablehlo.broadcast_in_dim %13937, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13939 = stablehlo.multiply %13938, %4808 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13940 = stablehlo.broadcast_in_dim %13910, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13941 = stablehlo.broadcast_in_dim %13940, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13942 = stablehlo.multiply %13941, %arg332 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13943 = stablehlo.add %13939, %13942 : tensor<1x2048x8192xf32> loc(#loc5263)
    %13944 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13945 = stablehlo.broadcast_in_dim %13944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13946 = stablehlo.subtract %13945, %13933 : tensor<1xf32> loc(#loc5260)
    %13947 = stablehlo.multiply %4808, %4808 : tensor<1x2048x8192xf32> loc(#loc5264)
    %13948 = stablehlo.broadcast_in_dim %13946, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13949 = stablehlo.broadcast_in_dim %13948, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13950 = stablehlo.multiply %13949, %13947 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13951 = stablehlo.broadcast_in_dim %13933, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %13952 = stablehlo.broadcast_in_dim %13951, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %13953 = stablehlo.multiply %13952, %arg508 : tensor<1x2048x8192xf32> loc(#loc5262)
    %13954 = stablehlo.add %13950, %13953 : tensor<1x2048x8192xf32> loc(#loc5263)
    %13955 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13956 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13957 = stablehlo.broadcast_in_dim %13956, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13958 = stablehlo.add %13955, %13957 : tensor<1xf32> loc(#loc5259)
    %13959 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13960 = stablehlo.broadcast_in_dim %13959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13961 = stablehlo.subtract %13958, %13960 : tensor<1xf32> loc(#loc5260)
    %13962 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13963 = stablehlo.broadcast_in_dim %13962, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13964 = stablehlo.power %13963, %13961 : tensor<1xf32> loc(#loc5261)
    %13965 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13966 = stablehlo.broadcast_in_dim %13965, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13967 = stablehlo.subtract %13966, %13964 : tensor<1xf32> loc(#loc5260)
    %13968 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13969 = stablehlo.broadcast_in_dim %13968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13970 = stablehlo.multiply %13969, %13967 : tensor<1xf32> loc(#loc5236)
    %13971 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %13972 = stablehlo.broadcast_in_dim %13971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13973 = stablehlo.power %13972, %13958 : tensor<1xf32> loc(#loc5261)
    %13974 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13975 = stablehlo.broadcast_in_dim %13974, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13976 = stablehlo.subtract %13975, %13973 : tensor<1xf32> loc(#loc5260)
    %13977 = stablehlo.divide %13970, %13976 : tensor<1xf32> loc(#loc5253)
    %13978 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %13979 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13980 = stablehlo.broadcast_in_dim %13979, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %13981 = stablehlo.add %13978, %13980 : tensor<1xf32> loc(#loc5259)
    %13982 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13983 = stablehlo.broadcast_in_dim %13982, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13984 = stablehlo.subtract %13981, %13983 : tensor<1xf32> loc(#loc5260)
    %13985 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13986 = stablehlo.broadcast_in_dim %13985, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13987 = stablehlo.power %13986, %13984 : tensor<1xf32> loc(#loc5261)
    %13988 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13989 = stablehlo.broadcast_in_dim %13988, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13990 = stablehlo.subtract %13989, %13987 : tensor<1xf32> loc(#loc5260)
    %13991 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13992 = stablehlo.broadcast_in_dim %13991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %13993 = stablehlo.multiply %13992, %13990 : tensor<1xf32> loc(#loc5236)
    %13994 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %13995 = stablehlo.broadcast_in_dim %13994, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %13996 = stablehlo.power %13995, %13981 : tensor<1xf32> loc(#loc5261)
    %13997 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %13998 = stablehlo.broadcast_in_dim %13997, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %13999 = stablehlo.subtract %13998, %13996 : tensor<1xf32> loc(#loc5260)
    %14000 = stablehlo.divide %13993, %13999 : tensor<1xf32> loc(#loc5253)
    %14001 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14002 = stablehlo.broadcast_in_dim %14001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14003 = stablehlo.subtract %14002, %13977 : tensor<1xf32> loc(#loc5260)
    %14004 = stablehlo.broadcast_in_dim %14003, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14005 = stablehlo.broadcast_in_dim %14004, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14006 = stablehlo.multiply %14005, %4812 : tensor<1x2048xf32> loc(#loc5262)
    %14007 = stablehlo.broadcast_in_dim %13977, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14008 = stablehlo.broadcast_in_dim %14007, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14009 = stablehlo.multiply %14008, %arg333 : tensor<1x2048xf32> loc(#loc5262)
    %14010 = stablehlo.add %14006, %14009 : tensor<1x2048xf32> loc(#loc5263)
    %14011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14012 = stablehlo.broadcast_in_dim %14011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14013 = stablehlo.subtract %14012, %14000 : tensor<1xf32> loc(#loc5260)
    %14014 = stablehlo.multiply %4812, %4812 : tensor<1x2048xf32> loc(#loc5264)
    %14015 = stablehlo.broadcast_in_dim %14013, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14016 = stablehlo.broadcast_in_dim %14015, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14017 = stablehlo.multiply %14016, %14014 : tensor<1x2048xf32> loc(#loc5262)
    %14018 = stablehlo.broadcast_in_dim %14000, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14019 = stablehlo.broadcast_in_dim %14018, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14020 = stablehlo.multiply %14019, %arg509 : tensor<1x2048xf32> loc(#loc5262)
    %14021 = stablehlo.add %14017, %14020 : tensor<1x2048xf32> loc(#loc5263)
    %14022 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14023 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14024 = stablehlo.broadcast_in_dim %14023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14025 = stablehlo.add %14022, %14024 : tensor<1xf32> loc(#loc5259)
    %14026 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14027 = stablehlo.broadcast_in_dim %14026, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14028 = stablehlo.subtract %14025, %14027 : tensor<1xf32> loc(#loc5260)
    %14029 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14030 = stablehlo.broadcast_in_dim %14029, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14031 = stablehlo.power %14030, %14028 : tensor<1xf32> loc(#loc5261)
    %14032 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14033 = stablehlo.broadcast_in_dim %14032, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14034 = stablehlo.subtract %14033, %14031 : tensor<1xf32> loc(#loc5260)
    %14035 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14036 = stablehlo.broadcast_in_dim %14035, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14037 = stablehlo.multiply %14036, %14034 : tensor<1xf32> loc(#loc5236)
    %14038 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14039 = stablehlo.broadcast_in_dim %14038, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14040 = stablehlo.power %14039, %14025 : tensor<1xf32> loc(#loc5261)
    %14041 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14042 = stablehlo.broadcast_in_dim %14041, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14043 = stablehlo.subtract %14042, %14040 : tensor<1xf32> loc(#loc5260)
    %14044 = stablehlo.divide %14037, %14043 : tensor<1xf32> loc(#loc5253)
    %14045 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14046 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14047 = stablehlo.broadcast_in_dim %14046, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14048 = stablehlo.add %14045, %14047 : tensor<1xf32> loc(#loc5259)
    %14049 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14050 = stablehlo.broadcast_in_dim %14049, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14051 = stablehlo.subtract %14048, %14050 : tensor<1xf32> loc(#loc5260)
    %14052 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14053 = stablehlo.broadcast_in_dim %14052, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14054 = stablehlo.power %14053, %14051 : tensor<1xf32> loc(#loc5261)
    %14055 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14056 = stablehlo.broadcast_in_dim %14055, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14057 = stablehlo.subtract %14056, %14054 : tensor<1xf32> loc(#loc5260)
    %14058 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14059 = stablehlo.broadcast_in_dim %14058, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14060 = stablehlo.multiply %14059, %14057 : tensor<1xf32> loc(#loc5236)
    %14061 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14062 = stablehlo.broadcast_in_dim %14061, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14063 = stablehlo.power %14062, %14048 : tensor<1xf32> loc(#loc5261)
    %14064 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14065 = stablehlo.broadcast_in_dim %14064, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14066 = stablehlo.subtract %14065, %14063 : tensor<1xf32> loc(#loc5260)
    %14067 = stablehlo.divide %14060, %14066 : tensor<1xf32> loc(#loc5253)
    %14068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14069 = stablehlo.broadcast_in_dim %14068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14070 = stablehlo.subtract %14069, %14044 : tensor<1xf32> loc(#loc5260)
    %14071 = stablehlo.broadcast_in_dim %14070, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14072 = stablehlo.broadcast_in_dim %14071, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14073 = stablehlo.multiply %14072, %4816 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14074 = stablehlo.broadcast_in_dim %14044, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14075 = stablehlo.broadcast_in_dim %14074, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14076 = stablehlo.multiply %14075, %arg334 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14077 = stablehlo.add %14073, %14076 : tensor<1x8192x2048xf32> loc(#loc5263)
    %14078 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14079 = stablehlo.broadcast_in_dim %14078, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14080 = stablehlo.subtract %14079, %14067 : tensor<1xf32> loc(#loc5260)
    %14081 = stablehlo.multiply %4816, %4816 : tensor<1x8192x2048xf32> loc(#loc5264)
    %14082 = stablehlo.broadcast_in_dim %14080, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14083 = stablehlo.broadcast_in_dim %14082, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14084 = stablehlo.multiply %14083, %14081 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14085 = stablehlo.broadcast_in_dim %14067, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14086 = stablehlo.broadcast_in_dim %14085, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14087 = stablehlo.multiply %14086, %arg510 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14088 = stablehlo.add %14084, %14087 : tensor<1x8192x2048xf32> loc(#loc5263)
    %14089 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14090 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14091 = stablehlo.broadcast_in_dim %14090, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14092 = stablehlo.add %14089, %14091 : tensor<1xf32> loc(#loc5259)
    %14093 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14094 = stablehlo.broadcast_in_dim %14093, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14095 = stablehlo.subtract %14092, %14094 : tensor<1xf32> loc(#loc5260)
    %14096 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14097 = stablehlo.broadcast_in_dim %14096, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14098 = stablehlo.power %14097, %14095 : tensor<1xf32> loc(#loc5261)
    %14099 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14100 = stablehlo.broadcast_in_dim %14099, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14101 = stablehlo.subtract %14100, %14098 : tensor<1xf32> loc(#loc5260)
    %14102 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14103 = stablehlo.broadcast_in_dim %14102, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14104 = stablehlo.multiply %14103, %14101 : tensor<1xf32> loc(#loc5236)
    %14105 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14106 = stablehlo.broadcast_in_dim %14105, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14107 = stablehlo.power %14106, %14092 : tensor<1xf32> loc(#loc5261)
    %14108 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14109 = stablehlo.broadcast_in_dim %14108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14110 = stablehlo.subtract %14109, %14107 : tensor<1xf32> loc(#loc5260)
    %14111 = stablehlo.divide %14104, %14110 : tensor<1xf32> loc(#loc5253)
    %14112 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14113 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14114 = stablehlo.broadcast_in_dim %14113, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14115 = stablehlo.add %14112, %14114 : tensor<1xf32> loc(#loc5259)
    %14116 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14117 = stablehlo.broadcast_in_dim %14116, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14118 = stablehlo.subtract %14115, %14117 : tensor<1xf32> loc(#loc5260)
    %14119 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14120 = stablehlo.broadcast_in_dim %14119, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14121 = stablehlo.power %14120, %14118 : tensor<1xf32> loc(#loc5261)
    %14122 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14123 = stablehlo.broadcast_in_dim %14122, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14124 = stablehlo.subtract %14123, %14121 : tensor<1xf32> loc(#loc5260)
    %14125 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14126 = stablehlo.broadcast_in_dim %14125, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14127 = stablehlo.multiply %14126, %14124 : tensor<1xf32> loc(#loc5236)
    %14128 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14129 = stablehlo.broadcast_in_dim %14128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14130 = stablehlo.power %14129, %14115 : tensor<1xf32> loc(#loc5261)
    %14131 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14132 = stablehlo.broadcast_in_dim %14131, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14133 = stablehlo.subtract %14132, %14130 : tensor<1xf32> loc(#loc5260)
    %14134 = stablehlo.divide %14127, %14133 : tensor<1xf32> loc(#loc5253)
    %14135 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14136 = stablehlo.broadcast_in_dim %14135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14137 = stablehlo.subtract %14136, %14111 : tensor<1xf32> loc(#loc5260)
    %14138 = stablehlo.broadcast_in_dim %14137, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14139 = stablehlo.broadcast_in_dim %14138, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14140 = stablehlo.multiply %14139, %4820 : tensor<1x2048xf32> loc(#loc5262)
    %14141 = stablehlo.broadcast_in_dim %14111, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14142 = stablehlo.broadcast_in_dim %14141, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14143 = stablehlo.multiply %14142, %arg335 : tensor<1x2048xf32> loc(#loc5262)
    %14144 = stablehlo.add %14140, %14143 : tensor<1x2048xf32> loc(#loc5263)
    %14145 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14146 = stablehlo.broadcast_in_dim %14145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14147 = stablehlo.subtract %14146, %14134 : tensor<1xf32> loc(#loc5260)
    %14148 = stablehlo.multiply %4820, %4820 : tensor<1x2048xf32> loc(#loc5264)
    %14149 = stablehlo.broadcast_in_dim %14147, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14150 = stablehlo.broadcast_in_dim %14149, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14151 = stablehlo.multiply %14150, %14148 : tensor<1x2048xf32> loc(#loc5262)
    %14152 = stablehlo.broadcast_in_dim %14134, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14153 = stablehlo.broadcast_in_dim %14152, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14154 = stablehlo.multiply %14153, %arg511 : tensor<1x2048xf32> loc(#loc5262)
    %14155 = stablehlo.add %14151, %14154 : tensor<1x2048xf32> loc(#loc5263)
    %14156 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14157 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14158 = stablehlo.broadcast_in_dim %14157, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14159 = stablehlo.add %14156, %14158 : tensor<1xf32> loc(#loc5259)
    %14160 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14161 = stablehlo.broadcast_in_dim %14160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14162 = stablehlo.subtract %14159, %14161 : tensor<1xf32> loc(#loc5260)
    %14163 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14164 = stablehlo.broadcast_in_dim %14163, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14165 = stablehlo.power %14164, %14162 : tensor<1xf32> loc(#loc5261)
    %14166 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14167 = stablehlo.broadcast_in_dim %14166, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14168 = stablehlo.subtract %14167, %14165 : tensor<1xf32> loc(#loc5260)
    %14169 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14170 = stablehlo.broadcast_in_dim %14169, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14171 = stablehlo.multiply %14170, %14168 : tensor<1xf32> loc(#loc5236)
    %14172 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14173 = stablehlo.broadcast_in_dim %14172, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14174 = stablehlo.power %14173, %14159 : tensor<1xf32> loc(#loc5261)
    %14175 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14176 = stablehlo.broadcast_in_dim %14175, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14177 = stablehlo.subtract %14176, %14174 : tensor<1xf32> loc(#loc5260)
    %14178 = stablehlo.divide %14171, %14177 : tensor<1xf32> loc(#loc5253)
    %14179 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14180 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14181 = stablehlo.broadcast_in_dim %14180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14182 = stablehlo.add %14179, %14181 : tensor<1xf32> loc(#loc5259)
    %14183 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14184 = stablehlo.broadcast_in_dim %14183, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14185 = stablehlo.subtract %14182, %14184 : tensor<1xf32> loc(#loc5260)
    %14186 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14187 = stablehlo.broadcast_in_dim %14186, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14188 = stablehlo.power %14187, %14185 : tensor<1xf32> loc(#loc5261)
    %14189 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14190 = stablehlo.broadcast_in_dim %14189, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14191 = stablehlo.subtract %14190, %14188 : tensor<1xf32> loc(#loc5260)
    %14192 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14193 = stablehlo.broadcast_in_dim %14192, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14194 = stablehlo.multiply %14193, %14191 : tensor<1xf32> loc(#loc5236)
    %14195 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14196 = stablehlo.broadcast_in_dim %14195, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14197 = stablehlo.power %14196, %14182 : tensor<1xf32> loc(#loc5261)
    %14198 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14199 = stablehlo.broadcast_in_dim %14198, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14200 = stablehlo.subtract %14199, %14197 : tensor<1xf32> loc(#loc5260)
    %14201 = stablehlo.divide %14194, %14200 : tensor<1xf32> loc(#loc5253)
    %14202 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14203 = stablehlo.broadcast_in_dim %14202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14204 = stablehlo.subtract %14203, %14178 : tensor<1xf32> loc(#loc5260)
    %14205 = stablehlo.broadcast_in_dim %14204, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14206 = stablehlo.broadcast_in_dim %14205, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14207 = stablehlo.multiply %14206, %4824 : tensor<1x2048xf32> loc(#loc5262)
    %14208 = stablehlo.broadcast_in_dim %14178, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14209 = stablehlo.broadcast_in_dim %14208, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14210 = stablehlo.multiply %14209, %arg336 : tensor<1x2048xf32> loc(#loc5262)
    %14211 = stablehlo.add %14207, %14210 : tensor<1x2048xf32> loc(#loc5263)
    %14212 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14213 = stablehlo.broadcast_in_dim %14212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14214 = stablehlo.subtract %14213, %14201 : tensor<1xf32> loc(#loc5260)
    %14215 = stablehlo.multiply %4824, %4824 : tensor<1x2048xf32> loc(#loc5264)
    %14216 = stablehlo.broadcast_in_dim %14214, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14217 = stablehlo.broadcast_in_dim %14216, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14218 = stablehlo.multiply %14217, %14215 : tensor<1x2048xf32> loc(#loc5262)
    %14219 = stablehlo.broadcast_in_dim %14201, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14220 = stablehlo.broadcast_in_dim %14219, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14221 = stablehlo.multiply %14220, %arg512 : tensor<1x2048xf32> loc(#loc5262)
    %14222 = stablehlo.add %14218, %14221 : tensor<1x2048xf32> loc(#loc5263)
    %14223 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14224 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14225 = stablehlo.broadcast_in_dim %14224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14226 = stablehlo.add %14223, %14225 : tensor<1xf32> loc(#loc5259)
    %14227 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14228 = stablehlo.broadcast_in_dim %14227, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14229 = stablehlo.subtract %14226, %14228 : tensor<1xf32> loc(#loc5260)
    %14230 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14231 = stablehlo.broadcast_in_dim %14230, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14232 = stablehlo.power %14231, %14229 : tensor<1xf32> loc(#loc5261)
    %14233 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14234 = stablehlo.broadcast_in_dim %14233, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14235 = stablehlo.subtract %14234, %14232 : tensor<1xf32> loc(#loc5260)
    %14236 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14237 = stablehlo.broadcast_in_dim %14236, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14238 = stablehlo.multiply %14237, %14235 : tensor<1xf32> loc(#loc5236)
    %14239 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14240 = stablehlo.broadcast_in_dim %14239, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14241 = stablehlo.power %14240, %14226 : tensor<1xf32> loc(#loc5261)
    %14242 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14243 = stablehlo.broadcast_in_dim %14242, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14244 = stablehlo.subtract %14243, %14241 : tensor<1xf32> loc(#loc5260)
    %14245 = stablehlo.divide %14238, %14244 : tensor<1xf32> loc(#loc5253)
    %14246 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14247 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14248 = stablehlo.broadcast_in_dim %14247, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14249 = stablehlo.add %14246, %14248 : tensor<1xf32> loc(#loc5259)
    %14250 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14251 = stablehlo.broadcast_in_dim %14250, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14252 = stablehlo.subtract %14249, %14251 : tensor<1xf32> loc(#loc5260)
    %14253 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14254 = stablehlo.broadcast_in_dim %14253, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14255 = stablehlo.power %14254, %14252 : tensor<1xf32> loc(#loc5261)
    %14256 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14257 = stablehlo.broadcast_in_dim %14256, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14258 = stablehlo.subtract %14257, %14255 : tensor<1xf32> loc(#loc5260)
    %14259 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14260 = stablehlo.broadcast_in_dim %14259, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14261 = stablehlo.multiply %14260, %14258 : tensor<1xf32> loc(#loc5236)
    %14262 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14263 = stablehlo.broadcast_in_dim %14262, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14264 = stablehlo.power %14263, %14249 : tensor<1xf32> loc(#loc5261)
    %14265 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14266 = stablehlo.broadcast_in_dim %14265, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14267 = stablehlo.subtract %14266, %14264 : tensor<1xf32> loc(#loc5260)
    %14268 = stablehlo.divide %14261, %14267 : tensor<1xf32> loc(#loc5253)
    %14269 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14270 = stablehlo.broadcast_in_dim %14269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14271 = stablehlo.subtract %14270, %14245 : tensor<1xf32> loc(#loc5260)
    %14272 = stablehlo.broadcast_in_dim %14271, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14273 = stablehlo.broadcast_in_dim %14272, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14274 = stablehlo.multiply %14273, %4828 : tensor<1x2048xf32> loc(#loc5262)
    %14275 = stablehlo.broadcast_in_dim %14245, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14276 = stablehlo.broadcast_in_dim %14275, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14277 = stablehlo.multiply %14276, %arg337 : tensor<1x2048xf32> loc(#loc5262)
    %14278 = stablehlo.add %14274, %14277 : tensor<1x2048xf32> loc(#loc5263)
    %14279 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14280 = stablehlo.broadcast_in_dim %14279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14281 = stablehlo.subtract %14280, %14268 : tensor<1xf32> loc(#loc5260)
    %14282 = stablehlo.multiply %4828, %4828 : tensor<1x2048xf32> loc(#loc5264)
    %14283 = stablehlo.broadcast_in_dim %14281, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14284 = stablehlo.broadcast_in_dim %14283, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14285 = stablehlo.multiply %14284, %14282 : tensor<1x2048xf32> loc(#loc5262)
    %14286 = stablehlo.broadcast_in_dim %14268, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14287 = stablehlo.broadcast_in_dim %14286, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14288 = stablehlo.multiply %14287, %arg513 : tensor<1x2048xf32> loc(#loc5262)
    %14289 = stablehlo.add %14285, %14288 : tensor<1x2048xf32> loc(#loc5263)
    %14290 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14291 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14292 = stablehlo.broadcast_in_dim %14291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14293 = stablehlo.add %14290, %14292 : tensor<1xf32> loc(#loc5259)
    %14294 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14295 = stablehlo.broadcast_in_dim %14294, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14296 = stablehlo.subtract %14293, %14295 : tensor<1xf32> loc(#loc5260)
    %14297 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14298 = stablehlo.broadcast_in_dim %14297, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14299 = stablehlo.power %14298, %14296 : tensor<1xf32> loc(#loc5261)
    %14300 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14301 = stablehlo.broadcast_in_dim %14300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14302 = stablehlo.subtract %14301, %14299 : tensor<1xf32> loc(#loc5260)
    %14303 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14304 = stablehlo.broadcast_in_dim %14303, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14305 = stablehlo.multiply %14304, %14302 : tensor<1xf32> loc(#loc5236)
    %14306 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14307 = stablehlo.broadcast_in_dim %14306, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14308 = stablehlo.power %14307, %14293 : tensor<1xf32> loc(#loc5261)
    %14309 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14310 = stablehlo.broadcast_in_dim %14309, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14311 = stablehlo.subtract %14310, %14308 : tensor<1xf32> loc(#loc5260)
    %14312 = stablehlo.divide %14305, %14311 : tensor<1xf32> loc(#loc5253)
    %14313 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14314 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14315 = stablehlo.broadcast_in_dim %14314, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14316 = stablehlo.add %14313, %14315 : tensor<1xf32> loc(#loc5259)
    %14317 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14318 = stablehlo.broadcast_in_dim %14317, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14319 = stablehlo.subtract %14316, %14318 : tensor<1xf32> loc(#loc5260)
    %14320 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14321 = stablehlo.broadcast_in_dim %14320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14322 = stablehlo.power %14321, %14319 : tensor<1xf32> loc(#loc5261)
    %14323 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14324 = stablehlo.broadcast_in_dim %14323, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14325 = stablehlo.subtract %14324, %14322 : tensor<1xf32> loc(#loc5260)
    %14326 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14327 = stablehlo.broadcast_in_dim %14326, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14328 = stablehlo.multiply %14327, %14325 : tensor<1xf32> loc(#loc5236)
    %14329 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14330 = stablehlo.broadcast_in_dim %14329, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14331 = stablehlo.power %14330, %14316 : tensor<1xf32> loc(#loc5261)
    %14332 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14333 = stablehlo.broadcast_in_dim %14332, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14334 = stablehlo.subtract %14333, %14331 : tensor<1xf32> loc(#loc5260)
    %14335 = stablehlo.divide %14328, %14334 : tensor<1xf32> loc(#loc5253)
    %14336 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14337 = stablehlo.broadcast_in_dim %14336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14338 = stablehlo.subtract %14337, %14312 : tensor<1xf32> loc(#loc5260)
    %14339 = stablehlo.broadcast_in_dim %14338, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14340 = stablehlo.broadcast_in_dim %14339, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14341 = stablehlo.multiply %14340, %4832 : tensor<1x2048xf32> loc(#loc5262)
    %14342 = stablehlo.broadcast_in_dim %14312, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14343 = stablehlo.broadcast_in_dim %14342, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14344 = stablehlo.multiply %14343, %arg338 : tensor<1x2048xf32> loc(#loc5262)
    %14345 = stablehlo.add %14341, %14344 : tensor<1x2048xf32> loc(#loc5263)
    %14346 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14347 = stablehlo.broadcast_in_dim %14346, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14348 = stablehlo.subtract %14347, %14335 : tensor<1xf32> loc(#loc5260)
    %14349 = stablehlo.multiply %4832, %4832 : tensor<1x2048xf32> loc(#loc5264)
    %14350 = stablehlo.broadcast_in_dim %14348, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14351 = stablehlo.broadcast_in_dim %14350, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14352 = stablehlo.multiply %14351, %14349 : tensor<1x2048xf32> loc(#loc5262)
    %14353 = stablehlo.broadcast_in_dim %14335, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14354 = stablehlo.broadcast_in_dim %14353, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14355 = stablehlo.multiply %14354, %arg514 : tensor<1x2048xf32> loc(#loc5262)
    %14356 = stablehlo.add %14352, %14355 : tensor<1x2048xf32> loc(#loc5263)
    %14357 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14358 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14359 = stablehlo.broadcast_in_dim %14358, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14360 = stablehlo.add %14357, %14359 : tensor<1xf32> loc(#loc5259)
    %14361 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14362 = stablehlo.broadcast_in_dim %14361, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14363 = stablehlo.subtract %14360, %14362 : tensor<1xf32> loc(#loc5260)
    %14364 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14365 = stablehlo.broadcast_in_dim %14364, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14366 = stablehlo.power %14365, %14363 : tensor<1xf32> loc(#loc5261)
    %14367 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14368 = stablehlo.broadcast_in_dim %14367, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14369 = stablehlo.subtract %14368, %14366 : tensor<1xf32> loc(#loc5260)
    %14370 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14371 = stablehlo.broadcast_in_dim %14370, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14372 = stablehlo.multiply %14371, %14369 : tensor<1xf32> loc(#loc5236)
    %14373 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14374 = stablehlo.broadcast_in_dim %14373, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14375 = stablehlo.power %14374, %14360 : tensor<1xf32> loc(#loc5261)
    %14376 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14377 = stablehlo.broadcast_in_dim %14376, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14378 = stablehlo.subtract %14377, %14375 : tensor<1xf32> loc(#loc5260)
    %14379 = stablehlo.divide %14372, %14378 : tensor<1xf32> loc(#loc5253)
    %14380 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14381 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14382 = stablehlo.broadcast_in_dim %14381, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14383 = stablehlo.add %14380, %14382 : tensor<1xf32> loc(#loc5259)
    %14384 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14385 = stablehlo.broadcast_in_dim %14384, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14386 = stablehlo.subtract %14383, %14385 : tensor<1xf32> loc(#loc5260)
    %14387 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14388 = stablehlo.broadcast_in_dim %14387, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14389 = stablehlo.power %14388, %14386 : tensor<1xf32> loc(#loc5261)
    %14390 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14391 = stablehlo.broadcast_in_dim %14390, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14392 = stablehlo.subtract %14391, %14389 : tensor<1xf32> loc(#loc5260)
    %14393 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14394 = stablehlo.broadcast_in_dim %14393, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14395 = stablehlo.multiply %14394, %14392 : tensor<1xf32> loc(#loc5236)
    %14396 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14397 = stablehlo.broadcast_in_dim %14396, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14398 = stablehlo.power %14397, %14383 : tensor<1xf32> loc(#loc5261)
    %14399 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14400 = stablehlo.broadcast_in_dim %14399, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14401 = stablehlo.subtract %14400, %14398 : tensor<1xf32> loc(#loc5260)
    %14402 = stablehlo.divide %14395, %14401 : tensor<1xf32> loc(#loc5253)
    %14403 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14404 = stablehlo.broadcast_in_dim %14403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14405 = stablehlo.subtract %14404, %14379 : tensor<1xf32> loc(#loc5260)
    %14406 = stablehlo.broadcast_in_dim %14405, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %14407 = stablehlo.broadcast_in_dim %14406, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14408 = stablehlo.multiply %14407, %4836 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14409 = stablehlo.broadcast_in_dim %14379, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %14410 = stablehlo.broadcast_in_dim %14409, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14411 = stablehlo.multiply %14410, %arg339 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14412 = stablehlo.add %14408, %14411 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %14413 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14414 = stablehlo.broadcast_in_dim %14413, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14415 = stablehlo.subtract %14414, %14402 : tensor<1xf32> loc(#loc5260)
    %14416 = stablehlo.multiply %4836, %4836 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %14417 = stablehlo.broadcast_in_dim %14415, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %14418 = stablehlo.broadcast_in_dim %14417, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14419 = stablehlo.multiply %14418, %14416 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14420 = stablehlo.broadcast_in_dim %14402, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %14421 = stablehlo.broadcast_in_dim %14420, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14422 = stablehlo.multiply %14421, %arg515 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %14423 = stablehlo.add %14419, %14422 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %14424 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14425 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14426 = stablehlo.broadcast_in_dim %14425, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14427 = stablehlo.add %14424, %14426 : tensor<1xf32> loc(#loc5259)
    %14428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14429 = stablehlo.broadcast_in_dim %14428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14430 = stablehlo.subtract %14427, %14429 : tensor<1xf32> loc(#loc5260)
    %14431 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14432 = stablehlo.broadcast_in_dim %14431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14433 = stablehlo.power %14432, %14430 : tensor<1xf32> loc(#loc5261)
    %14434 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14435 = stablehlo.broadcast_in_dim %14434, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14436 = stablehlo.subtract %14435, %14433 : tensor<1xf32> loc(#loc5260)
    %14437 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14438 = stablehlo.broadcast_in_dim %14437, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14439 = stablehlo.multiply %14438, %14436 : tensor<1xf32> loc(#loc5236)
    %14440 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14441 = stablehlo.broadcast_in_dim %14440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14442 = stablehlo.power %14441, %14427 : tensor<1xf32> loc(#loc5261)
    %14443 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14444 = stablehlo.broadcast_in_dim %14443, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14445 = stablehlo.subtract %14444, %14442 : tensor<1xf32> loc(#loc5260)
    %14446 = stablehlo.divide %14439, %14445 : tensor<1xf32> loc(#loc5253)
    %14447 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14448 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14449 = stablehlo.broadcast_in_dim %14448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14450 = stablehlo.add %14447, %14449 : tensor<1xf32> loc(#loc5259)
    %14451 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14452 = stablehlo.broadcast_in_dim %14451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14453 = stablehlo.subtract %14450, %14452 : tensor<1xf32> loc(#loc5260)
    %14454 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14455 = stablehlo.broadcast_in_dim %14454, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14456 = stablehlo.power %14455, %14453 : tensor<1xf32> loc(#loc5261)
    %14457 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14458 = stablehlo.broadcast_in_dim %14457, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14459 = stablehlo.subtract %14458, %14456 : tensor<1xf32> loc(#loc5260)
    %14460 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14461 = stablehlo.broadcast_in_dim %14460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14462 = stablehlo.multiply %14461, %14459 : tensor<1xf32> loc(#loc5236)
    %14463 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14464 = stablehlo.broadcast_in_dim %14463, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14465 = stablehlo.power %14464, %14450 : tensor<1xf32> loc(#loc5261)
    %14466 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14467 = stablehlo.broadcast_in_dim %14466, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14468 = stablehlo.subtract %14467, %14465 : tensor<1xf32> loc(#loc5260)
    %14469 = stablehlo.divide %14462, %14468 : tensor<1xf32> loc(#loc5253)
    %14470 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14471 = stablehlo.broadcast_in_dim %14470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14472 = stablehlo.subtract %14471, %14446 : tensor<1xf32> loc(#loc5260)
    %14473 = stablehlo.broadcast_in_dim %14472, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14474 = stablehlo.broadcast_in_dim %14473, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %14475 = stablehlo.multiply %14474, %4840 : tensor<1x64xf32> loc(#loc5262)
    %14476 = stablehlo.broadcast_in_dim %14446, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14477 = stablehlo.broadcast_in_dim %14476, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %14478 = stablehlo.multiply %14477, %arg340 : tensor<1x64xf32> loc(#loc5262)
    %14479 = stablehlo.add %14475, %14478 : tensor<1x64xf32> loc(#loc5263)
    %14480 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14481 = stablehlo.broadcast_in_dim %14480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14482 = stablehlo.subtract %14481, %14469 : tensor<1xf32> loc(#loc5260)
    %14483 = stablehlo.multiply %4840, %4840 : tensor<1x64xf32> loc(#loc5264)
    %14484 = stablehlo.broadcast_in_dim %14482, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14485 = stablehlo.broadcast_in_dim %14484, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %14486 = stablehlo.multiply %14485, %14483 : tensor<1x64xf32> loc(#loc5262)
    %14487 = stablehlo.broadcast_in_dim %14469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14488 = stablehlo.broadcast_in_dim %14487, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %14489 = stablehlo.multiply %14488, %arg516 : tensor<1x64xf32> loc(#loc5262)
    %14490 = stablehlo.add %14486, %14489 : tensor<1x64xf32> loc(#loc5263)
    %14491 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14492 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14493 = stablehlo.broadcast_in_dim %14492, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14494 = stablehlo.add %14491, %14493 : tensor<1xf32> loc(#loc5259)
    %14495 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14496 = stablehlo.broadcast_in_dim %14495, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14497 = stablehlo.subtract %14494, %14496 : tensor<1xf32> loc(#loc5260)
    %14498 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14499 = stablehlo.broadcast_in_dim %14498, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14500 = stablehlo.power %14499, %14497 : tensor<1xf32> loc(#loc5261)
    %14501 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14502 = stablehlo.broadcast_in_dim %14501, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14503 = stablehlo.subtract %14502, %14500 : tensor<1xf32> loc(#loc5260)
    %14504 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14505 = stablehlo.broadcast_in_dim %14504, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14506 = stablehlo.multiply %14505, %14503 : tensor<1xf32> loc(#loc5236)
    %14507 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14508 = stablehlo.broadcast_in_dim %14507, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14509 = stablehlo.power %14508, %14494 : tensor<1xf32> loc(#loc5261)
    %14510 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14511 = stablehlo.broadcast_in_dim %14510, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14512 = stablehlo.subtract %14511, %14509 : tensor<1xf32> loc(#loc5260)
    %14513 = stablehlo.divide %14506, %14512 : tensor<1xf32> loc(#loc5253)
    %14514 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14515 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14516 = stablehlo.broadcast_in_dim %14515, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14517 = stablehlo.add %14514, %14516 : tensor<1xf32> loc(#loc5259)
    %14518 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14519 = stablehlo.broadcast_in_dim %14518, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14520 = stablehlo.subtract %14517, %14519 : tensor<1xf32> loc(#loc5260)
    %14521 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14522 = stablehlo.broadcast_in_dim %14521, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14523 = stablehlo.power %14522, %14520 : tensor<1xf32> loc(#loc5261)
    %14524 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14525 = stablehlo.broadcast_in_dim %14524, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14526 = stablehlo.subtract %14525, %14523 : tensor<1xf32> loc(#loc5260)
    %14527 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14528 = stablehlo.broadcast_in_dim %14527, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14529 = stablehlo.multiply %14528, %14526 : tensor<1xf32> loc(#loc5236)
    %14530 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14531 = stablehlo.broadcast_in_dim %14530, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14532 = stablehlo.power %14531, %14517 : tensor<1xf32> loc(#loc5261)
    %14533 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14534 = stablehlo.broadcast_in_dim %14533, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14535 = stablehlo.subtract %14534, %14532 : tensor<1xf32> loc(#loc5260)
    %14536 = stablehlo.divide %14529, %14535 : tensor<1xf32> loc(#loc5253)
    %14537 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14538 = stablehlo.broadcast_in_dim %14537, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14539 = stablehlo.subtract %14538, %14513 : tensor<1xf32> loc(#loc5260)
    %14540 = stablehlo.broadcast_in_dim %14539, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %14541 = stablehlo.broadcast_in_dim %14540, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %14542 = stablehlo.multiply %14541, %4844 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %14543 = stablehlo.broadcast_in_dim %14513, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %14544 = stablehlo.broadcast_in_dim %14543, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %14545 = stablehlo.multiply %14544, %arg341 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %14546 = stablehlo.add %14542, %14545 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %14547 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14548 = stablehlo.broadcast_in_dim %14547, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14549 = stablehlo.subtract %14548, %14536 : tensor<1xf32> loc(#loc5260)
    %14550 = stablehlo.multiply %4844, %4844 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %14551 = stablehlo.broadcast_in_dim %14549, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %14552 = stablehlo.broadcast_in_dim %14551, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %14553 = stablehlo.multiply %14552, %14550 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %14554 = stablehlo.broadcast_in_dim %14536, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %14555 = stablehlo.broadcast_in_dim %14554, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %14556 = stablehlo.multiply %14555, %arg517 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %14557 = stablehlo.add %14553, %14556 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %14558 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14559 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14560 = stablehlo.broadcast_in_dim %14559, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14561 = stablehlo.add %14558, %14560 : tensor<1xf32> loc(#loc5259)
    %14562 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14563 = stablehlo.broadcast_in_dim %14562, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14564 = stablehlo.subtract %14561, %14563 : tensor<1xf32> loc(#loc5260)
    %14565 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14566 = stablehlo.broadcast_in_dim %14565, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14567 = stablehlo.power %14566, %14564 : tensor<1xf32> loc(#loc5261)
    %14568 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14569 = stablehlo.broadcast_in_dim %14568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14570 = stablehlo.subtract %14569, %14567 : tensor<1xf32> loc(#loc5260)
    %14571 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14572 = stablehlo.broadcast_in_dim %14571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14573 = stablehlo.multiply %14572, %14570 : tensor<1xf32> loc(#loc5236)
    %14574 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14575 = stablehlo.broadcast_in_dim %14574, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14576 = stablehlo.power %14575, %14561 : tensor<1xf32> loc(#loc5261)
    %14577 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14578 = stablehlo.broadcast_in_dim %14577, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14579 = stablehlo.subtract %14578, %14576 : tensor<1xf32> loc(#loc5260)
    %14580 = stablehlo.divide %14573, %14579 : tensor<1xf32> loc(#loc5253)
    %14581 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14582 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14583 = stablehlo.broadcast_in_dim %14582, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14584 = stablehlo.add %14581, %14583 : tensor<1xf32> loc(#loc5259)
    %14585 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14586 = stablehlo.broadcast_in_dim %14585, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14587 = stablehlo.subtract %14584, %14586 : tensor<1xf32> loc(#loc5260)
    %14588 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14589 = stablehlo.broadcast_in_dim %14588, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14590 = stablehlo.power %14589, %14587 : tensor<1xf32> loc(#loc5261)
    %14591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14592 = stablehlo.broadcast_in_dim %14591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14593 = stablehlo.subtract %14592, %14590 : tensor<1xf32> loc(#loc5260)
    %14594 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14595 = stablehlo.broadcast_in_dim %14594, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14596 = stablehlo.multiply %14595, %14593 : tensor<1xf32> loc(#loc5236)
    %14597 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14598 = stablehlo.broadcast_in_dim %14597, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14599 = stablehlo.power %14598, %14584 : tensor<1xf32> loc(#loc5261)
    %14600 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14601 = stablehlo.broadcast_in_dim %14600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14602 = stablehlo.subtract %14601, %14599 : tensor<1xf32> loc(#loc5260)
    %14603 = stablehlo.divide %14596, %14602 : tensor<1xf32> loc(#loc5253)
    %14604 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14605 = stablehlo.broadcast_in_dim %14604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14606 = stablehlo.subtract %14605, %14580 : tensor<1xf32> loc(#loc5260)
    %14607 = stablehlo.broadcast_in_dim %14606, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14608 = stablehlo.broadcast_in_dim %14607, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %14609 = stablehlo.multiply %14608, %4848 : tensor<1x8192xf32> loc(#loc5262)
    %14610 = stablehlo.broadcast_in_dim %14580, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14611 = stablehlo.broadcast_in_dim %14610, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %14612 = stablehlo.multiply %14611, %arg342 : tensor<1x8192xf32> loc(#loc5262)
    %14613 = stablehlo.add %14609, %14612 : tensor<1x8192xf32> loc(#loc5263)
    %14614 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14615 = stablehlo.broadcast_in_dim %14614, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14616 = stablehlo.subtract %14615, %14603 : tensor<1xf32> loc(#loc5260)
    %14617 = stablehlo.multiply %4848, %4848 : tensor<1x8192xf32> loc(#loc5264)
    %14618 = stablehlo.broadcast_in_dim %14616, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14619 = stablehlo.broadcast_in_dim %14618, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %14620 = stablehlo.multiply %14619, %14617 : tensor<1x8192xf32> loc(#loc5262)
    %14621 = stablehlo.broadcast_in_dim %14603, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14622 = stablehlo.broadcast_in_dim %14621, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %14623 = stablehlo.multiply %14622, %arg518 : tensor<1x8192xf32> loc(#loc5262)
    %14624 = stablehlo.add %14620, %14623 : tensor<1x8192xf32> loc(#loc5263)
    %14625 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14626 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14627 = stablehlo.broadcast_in_dim %14626, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14628 = stablehlo.add %14625, %14627 : tensor<1xf32> loc(#loc5259)
    %14629 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14630 = stablehlo.broadcast_in_dim %14629, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14631 = stablehlo.subtract %14628, %14630 : tensor<1xf32> loc(#loc5260)
    %14632 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14633 = stablehlo.broadcast_in_dim %14632, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14634 = stablehlo.power %14633, %14631 : tensor<1xf32> loc(#loc5261)
    %14635 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14636 = stablehlo.broadcast_in_dim %14635, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14637 = stablehlo.subtract %14636, %14634 : tensor<1xf32> loc(#loc5260)
    %14638 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14639 = stablehlo.broadcast_in_dim %14638, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14640 = stablehlo.multiply %14639, %14637 : tensor<1xf32> loc(#loc5236)
    %14641 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14642 = stablehlo.broadcast_in_dim %14641, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14643 = stablehlo.power %14642, %14628 : tensor<1xf32> loc(#loc5261)
    %14644 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14645 = stablehlo.broadcast_in_dim %14644, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14646 = stablehlo.subtract %14645, %14643 : tensor<1xf32> loc(#loc5260)
    %14647 = stablehlo.divide %14640, %14646 : tensor<1xf32> loc(#loc5253)
    %14648 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14649 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14650 = stablehlo.broadcast_in_dim %14649, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14651 = stablehlo.add %14648, %14650 : tensor<1xf32> loc(#loc5259)
    %14652 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14653 = stablehlo.broadcast_in_dim %14652, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14654 = stablehlo.subtract %14651, %14653 : tensor<1xf32> loc(#loc5260)
    %14655 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14656 = stablehlo.broadcast_in_dim %14655, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14657 = stablehlo.power %14656, %14654 : tensor<1xf32> loc(#loc5261)
    %14658 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14659 = stablehlo.broadcast_in_dim %14658, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14660 = stablehlo.subtract %14659, %14657 : tensor<1xf32> loc(#loc5260)
    %14661 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14662 = stablehlo.broadcast_in_dim %14661, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14663 = stablehlo.multiply %14662, %14660 : tensor<1xf32> loc(#loc5236)
    %14664 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14665 = stablehlo.broadcast_in_dim %14664, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14666 = stablehlo.power %14665, %14651 : tensor<1xf32> loc(#loc5261)
    %14667 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14668 = stablehlo.broadcast_in_dim %14667, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14669 = stablehlo.subtract %14668, %14666 : tensor<1xf32> loc(#loc5260)
    %14670 = stablehlo.divide %14663, %14669 : tensor<1xf32> loc(#loc5253)
    %14671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14672 = stablehlo.broadcast_in_dim %14671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14673 = stablehlo.subtract %14672, %14647 : tensor<1xf32> loc(#loc5260)
    %14674 = stablehlo.broadcast_in_dim %14673, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14675 = stablehlo.broadcast_in_dim %14674, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %14676 = stablehlo.multiply %14675, %4852 : tensor<1x2048x8192xf32> loc(#loc5262)
    %14677 = stablehlo.broadcast_in_dim %14647, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14678 = stablehlo.broadcast_in_dim %14677, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %14679 = stablehlo.multiply %14678, %arg343 : tensor<1x2048x8192xf32> loc(#loc5262)
    %14680 = stablehlo.add %14676, %14679 : tensor<1x2048x8192xf32> loc(#loc5263)
    %14681 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14682 = stablehlo.broadcast_in_dim %14681, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14683 = stablehlo.subtract %14682, %14670 : tensor<1xf32> loc(#loc5260)
    %14684 = stablehlo.multiply %4852, %4852 : tensor<1x2048x8192xf32> loc(#loc5264)
    %14685 = stablehlo.broadcast_in_dim %14683, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14686 = stablehlo.broadcast_in_dim %14685, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %14687 = stablehlo.multiply %14686, %14684 : tensor<1x2048x8192xf32> loc(#loc5262)
    %14688 = stablehlo.broadcast_in_dim %14670, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14689 = stablehlo.broadcast_in_dim %14688, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %14690 = stablehlo.multiply %14689, %arg519 : tensor<1x2048x8192xf32> loc(#loc5262)
    %14691 = stablehlo.add %14687, %14690 : tensor<1x2048x8192xf32> loc(#loc5263)
    %14692 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14693 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14694 = stablehlo.broadcast_in_dim %14693, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14695 = stablehlo.add %14692, %14694 : tensor<1xf32> loc(#loc5259)
    %14696 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14697 = stablehlo.broadcast_in_dim %14696, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14698 = stablehlo.subtract %14695, %14697 : tensor<1xf32> loc(#loc5260)
    %14699 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14700 = stablehlo.broadcast_in_dim %14699, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14701 = stablehlo.power %14700, %14698 : tensor<1xf32> loc(#loc5261)
    %14702 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14703 = stablehlo.broadcast_in_dim %14702, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14704 = stablehlo.subtract %14703, %14701 : tensor<1xf32> loc(#loc5260)
    %14705 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14706 = stablehlo.broadcast_in_dim %14705, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14707 = stablehlo.multiply %14706, %14704 : tensor<1xf32> loc(#loc5236)
    %14708 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14709 = stablehlo.broadcast_in_dim %14708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14710 = stablehlo.power %14709, %14695 : tensor<1xf32> loc(#loc5261)
    %14711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14712 = stablehlo.broadcast_in_dim %14711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14713 = stablehlo.subtract %14712, %14710 : tensor<1xf32> loc(#loc5260)
    %14714 = stablehlo.divide %14707, %14713 : tensor<1xf32> loc(#loc5253)
    %14715 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14716 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14717 = stablehlo.broadcast_in_dim %14716, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14718 = stablehlo.add %14715, %14717 : tensor<1xf32> loc(#loc5259)
    %14719 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14720 = stablehlo.broadcast_in_dim %14719, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14721 = stablehlo.subtract %14718, %14720 : tensor<1xf32> loc(#loc5260)
    %14722 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14723 = stablehlo.broadcast_in_dim %14722, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14724 = stablehlo.power %14723, %14721 : tensor<1xf32> loc(#loc5261)
    %14725 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14726 = stablehlo.broadcast_in_dim %14725, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14727 = stablehlo.subtract %14726, %14724 : tensor<1xf32> loc(#loc5260)
    %14728 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14729 = stablehlo.broadcast_in_dim %14728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14730 = stablehlo.multiply %14729, %14727 : tensor<1xf32> loc(#loc5236)
    %14731 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14732 = stablehlo.broadcast_in_dim %14731, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14733 = stablehlo.power %14732, %14718 : tensor<1xf32> loc(#loc5261)
    %14734 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14735 = stablehlo.broadcast_in_dim %14734, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14736 = stablehlo.subtract %14735, %14733 : tensor<1xf32> loc(#loc5260)
    %14737 = stablehlo.divide %14730, %14736 : tensor<1xf32> loc(#loc5253)
    %14738 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14739 = stablehlo.broadcast_in_dim %14738, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14740 = stablehlo.subtract %14739, %14714 : tensor<1xf32> loc(#loc5260)
    %14741 = stablehlo.broadcast_in_dim %14740, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14742 = stablehlo.broadcast_in_dim %14741, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14743 = stablehlo.multiply %14742, %4856 : tensor<1x2048xf32> loc(#loc5262)
    %14744 = stablehlo.broadcast_in_dim %14714, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14745 = stablehlo.broadcast_in_dim %14744, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14746 = stablehlo.multiply %14745, %arg344 : tensor<1x2048xf32> loc(#loc5262)
    %14747 = stablehlo.add %14743, %14746 : tensor<1x2048xf32> loc(#loc5263)
    %14748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14749 = stablehlo.broadcast_in_dim %14748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14750 = stablehlo.subtract %14749, %14737 : tensor<1xf32> loc(#loc5260)
    %14751 = stablehlo.multiply %4856, %4856 : tensor<1x2048xf32> loc(#loc5264)
    %14752 = stablehlo.broadcast_in_dim %14750, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14753 = stablehlo.broadcast_in_dim %14752, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14754 = stablehlo.multiply %14753, %14751 : tensor<1x2048xf32> loc(#loc5262)
    %14755 = stablehlo.broadcast_in_dim %14737, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14756 = stablehlo.broadcast_in_dim %14755, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14757 = stablehlo.multiply %14756, %arg520 : tensor<1x2048xf32> loc(#loc5262)
    %14758 = stablehlo.add %14754, %14757 : tensor<1x2048xf32> loc(#loc5263)
    %14759 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14760 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14761 = stablehlo.broadcast_in_dim %14760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14762 = stablehlo.add %14759, %14761 : tensor<1xf32> loc(#loc5259)
    %14763 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14764 = stablehlo.broadcast_in_dim %14763, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14765 = stablehlo.subtract %14762, %14764 : tensor<1xf32> loc(#loc5260)
    %14766 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14767 = stablehlo.broadcast_in_dim %14766, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14768 = stablehlo.power %14767, %14765 : tensor<1xf32> loc(#loc5261)
    %14769 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14770 = stablehlo.broadcast_in_dim %14769, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14771 = stablehlo.subtract %14770, %14768 : tensor<1xf32> loc(#loc5260)
    %14772 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14773 = stablehlo.broadcast_in_dim %14772, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14774 = stablehlo.multiply %14773, %14771 : tensor<1xf32> loc(#loc5236)
    %14775 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14776 = stablehlo.broadcast_in_dim %14775, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14777 = stablehlo.power %14776, %14762 : tensor<1xf32> loc(#loc5261)
    %14778 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14779 = stablehlo.broadcast_in_dim %14778, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14780 = stablehlo.subtract %14779, %14777 : tensor<1xf32> loc(#loc5260)
    %14781 = stablehlo.divide %14774, %14780 : tensor<1xf32> loc(#loc5253)
    %14782 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14783 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14784 = stablehlo.broadcast_in_dim %14783, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14785 = stablehlo.add %14782, %14784 : tensor<1xf32> loc(#loc5259)
    %14786 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14787 = stablehlo.broadcast_in_dim %14786, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14788 = stablehlo.subtract %14785, %14787 : tensor<1xf32> loc(#loc5260)
    %14789 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14790 = stablehlo.broadcast_in_dim %14789, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14791 = stablehlo.power %14790, %14788 : tensor<1xf32> loc(#loc5261)
    %14792 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14793 = stablehlo.broadcast_in_dim %14792, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14794 = stablehlo.subtract %14793, %14791 : tensor<1xf32> loc(#loc5260)
    %14795 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14796 = stablehlo.broadcast_in_dim %14795, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14797 = stablehlo.multiply %14796, %14794 : tensor<1xf32> loc(#loc5236)
    %14798 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14799 = stablehlo.broadcast_in_dim %14798, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14800 = stablehlo.power %14799, %14785 : tensor<1xf32> loc(#loc5261)
    %14801 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14802 = stablehlo.broadcast_in_dim %14801, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14803 = stablehlo.subtract %14802, %14800 : tensor<1xf32> loc(#loc5260)
    %14804 = stablehlo.divide %14797, %14803 : tensor<1xf32> loc(#loc5253)
    %14805 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14806 = stablehlo.broadcast_in_dim %14805, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14807 = stablehlo.subtract %14806, %14781 : tensor<1xf32> loc(#loc5260)
    %14808 = stablehlo.broadcast_in_dim %14807, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14809 = stablehlo.broadcast_in_dim %14808, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14810 = stablehlo.multiply %14809, %4860 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14811 = stablehlo.broadcast_in_dim %14781, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14812 = stablehlo.broadcast_in_dim %14811, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14813 = stablehlo.multiply %14812, %arg345 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14814 = stablehlo.add %14810, %14813 : tensor<1x8192x2048xf32> loc(#loc5263)
    %14815 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14816 = stablehlo.broadcast_in_dim %14815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14817 = stablehlo.subtract %14816, %14804 : tensor<1xf32> loc(#loc5260)
    %14818 = stablehlo.multiply %4860, %4860 : tensor<1x8192x2048xf32> loc(#loc5264)
    %14819 = stablehlo.broadcast_in_dim %14817, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14820 = stablehlo.broadcast_in_dim %14819, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14821 = stablehlo.multiply %14820, %14818 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14822 = stablehlo.broadcast_in_dim %14804, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %14823 = stablehlo.broadcast_in_dim %14822, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %14824 = stablehlo.multiply %14823, %arg521 : tensor<1x8192x2048xf32> loc(#loc5262)
    %14825 = stablehlo.add %14821, %14824 : tensor<1x8192x2048xf32> loc(#loc5263)
    %14826 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14827 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14828 = stablehlo.broadcast_in_dim %14827, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14829 = stablehlo.add %14826, %14828 : tensor<1xf32> loc(#loc5259)
    %14830 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14831 = stablehlo.broadcast_in_dim %14830, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14832 = stablehlo.subtract %14829, %14831 : tensor<1xf32> loc(#loc5260)
    %14833 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14834 = stablehlo.broadcast_in_dim %14833, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14835 = stablehlo.power %14834, %14832 : tensor<1xf32> loc(#loc5261)
    %14836 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14837 = stablehlo.broadcast_in_dim %14836, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14838 = stablehlo.subtract %14837, %14835 : tensor<1xf32> loc(#loc5260)
    %14839 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14840 = stablehlo.broadcast_in_dim %14839, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14841 = stablehlo.multiply %14840, %14838 : tensor<1xf32> loc(#loc5236)
    %14842 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14843 = stablehlo.broadcast_in_dim %14842, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14844 = stablehlo.power %14843, %14829 : tensor<1xf32> loc(#loc5261)
    %14845 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14846 = stablehlo.broadcast_in_dim %14845, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14847 = stablehlo.subtract %14846, %14844 : tensor<1xf32> loc(#loc5260)
    %14848 = stablehlo.divide %14841, %14847 : tensor<1xf32> loc(#loc5253)
    %14849 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14850 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14851 = stablehlo.broadcast_in_dim %14850, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14852 = stablehlo.add %14849, %14851 : tensor<1xf32> loc(#loc5259)
    %14853 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14854 = stablehlo.broadcast_in_dim %14853, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14855 = stablehlo.subtract %14852, %14854 : tensor<1xf32> loc(#loc5260)
    %14856 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14857 = stablehlo.broadcast_in_dim %14856, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14858 = stablehlo.power %14857, %14855 : tensor<1xf32> loc(#loc5261)
    %14859 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14860 = stablehlo.broadcast_in_dim %14859, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14861 = stablehlo.subtract %14860, %14858 : tensor<1xf32> loc(#loc5260)
    %14862 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14863 = stablehlo.broadcast_in_dim %14862, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14864 = stablehlo.multiply %14863, %14861 : tensor<1xf32> loc(#loc5236)
    %14865 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14866 = stablehlo.broadcast_in_dim %14865, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14867 = stablehlo.power %14866, %14852 : tensor<1xf32> loc(#loc5261)
    %14868 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14869 = stablehlo.broadcast_in_dim %14868, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14870 = stablehlo.subtract %14869, %14867 : tensor<1xf32> loc(#loc5260)
    %14871 = stablehlo.divide %14864, %14870 : tensor<1xf32> loc(#loc5253)
    %14872 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14873 = stablehlo.broadcast_in_dim %14872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14874 = stablehlo.subtract %14873, %14848 : tensor<1xf32> loc(#loc5260)
    %14875 = stablehlo.broadcast_in_dim %14874, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14876 = stablehlo.broadcast_in_dim %14875, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14877 = stablehlo.multiply %14876, %4864 : tensor<1x2048xf32> loc(#loc5262)
    %14878 = stablehlo.broadcast_in_dim %14848, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14879 = stablehlo.broadcast_in_dim %14878, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14880 = stablehlo.multiply %14879, %arg346 : tensor<1x2048xf32> loc(#loc5262)
    %14881 = stablehlo.add %14877, %14880 : tensor<1x2048xf32> loc(#loc5263)
    %14882 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14883 = stablehlo.broadcast_in_dim %14882, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14884 = stablehlo.subtract %14883, %14871 : tensor<1xf32> loc(#loc5260)
    %14885 = stablehlo.multiply %4864, %4864 : tensor<1x2048xf32> loc(#loc5264)
    %14886 = stablehlo.broadcast_in_dim %14884, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14887 = stablehlo.broadcast_in_dim %14886, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14888 = stablehlo.multiply %14887, %14885 : tensor<1x2048xf32> loc(#loc5262)
    %14889 = stablehlo.broadcast_in_dim %14871, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14890 = stablehlo.broadcast_in_dim %14889, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14891 = stablehlo.multiply %14890, %arg522 : tensor<1x2048xf32> loc(#loc5262)
    %14892 = stablehlo.add %14888, %14891 : tensor<1x2048xf32> loc(#loc5263)
    %14893 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14894 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14895 = stablehlo.broadcast_in_dim %14894, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14896 = stablehlo.add %14893, %14895 : tensor<1xf32> loc(#loc5259)
    %14897 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14898 = stablehlo.broadcast_in_dim %14897, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14899 = stablehlo.subtract %14896, %14898 : tensor<1xf32> loc(#loc5260)
    %14900 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14901 = stablehlo.broadcast_in_dim %14900, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14902 = stablehlo.power %14901, %14899 : tensor<1xf32> loc(#loc5261)
    %14903 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14904 = stablehlo.broadcast_in_dim %14903, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14905 = stablehlo.subtract %14904, %14902 : tensor<1xf32> loc(#loc5260)
    %14906 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14907 = stablehlo.broadcast_in_dim %14906, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14908 = stablehlo.multiply %14907, %14905 : tensor<1xf32> loc(#loc5236)
    %14909 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14910 = stablehlo.broadcast_in_dim %14909, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14911 = stablehlo.power %14910, %14896 : tensor<1xf32> loc(#loc5261)
    %14912 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14913 = stablehlo.broadcast_in_dim %14912, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14914 = stablehlo.subtract %14913, %14911 : tensor<1xf32> loc(#loc5260)
    %14915 = stablehlo.divide %14908, %14914 : tensor<1xf32> loc(#loc5253)
    %14916 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14917 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14918 = stablehlo.broadcast_in_dim %14917, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14919 = stablehlo.add %14916, %14918 : tensor<1xf32> loc(#loc5259)
    %14920 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14921 = stablehlo.broadcast_in_dim %14920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14922 = stablehlo.subtract %14919, %14921 : tensor<1xf32> loc(#loc5260)
    %14923 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14924 = stablehlo.broadcast_in_dim %14923, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14925 = stablehlo.power %14924, %14922 : tensor<1xf32> loc(#loc5261)
    %14926 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14927 = stablehlo.broadcast_in_dim %14926, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14928 = stablehlo.subtract %14927, %14925 : tensor<1xf32> loc(#loc5260)
    %14929 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14930 = stablehlo.broadcast_in_dim %14929, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14931 = stablehlo.multiply %14930, %14928 : tensor<1xf32> loc(#loc5236)
    %14932 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14933 = stablehlo.broadcast_in_dim %14932, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14934 = stablehlo.power %14933, %14919 : tensor<1xf32> loc(#loc5261)
    %14935 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14936 = stablehlo.broadcast_in_dim %14935, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14937 = stablehlo.subtract %14936, %14934 : tensor<1xf32> loc(#loc5260)
    %14938 = stablehlo.divide %14931, %14937 : tensor<1xf32> loc(#loc5253)
    %14939 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14940 = stablehlo.broadcast_in_dim %14939, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14941 = stablehlo.subtract %14940, %14915 : tensor<1xf32> loc(#loc5260)
    %14942 = stablehlo.broadcast_in_dim %14941, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14943 = stablehlo.broadcast_in_dim %14942, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14944 = stablehlo.multiply %14943, %4868 : tensor<1x2048xf32> loc(#loc5262)
    %14945 = stablehlo.broadcast_in_dim %14915, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14946 = stablehlo.broadcast_in_dim %14945, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14947 = stablehlo.multiply %14946, %arg347 : tensor<1x2048xf32> loc(#loc5262)
    %14948 = stablehlo.add %14944, %14947 : tensor<1x2048xf32> loc(#loc5263)
    %14949 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14950 = stablehlo.broadcast_in_dim %14949, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14951 = stablehlo.subtract %14950, %14938 : tensor<1xf32> loc(#loc5260)
    %14952 = stablehlo.multiply %4868, %4868 : tensor<1x2048xf32> loc(#loc5264)
    %14953 = stablehlo.broadcast_in_dim %14951, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14954 = stablehlo.broadcast_in_dim %14953, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14955 = stablehlo.multiply %14954, %14952 : tensor<1x2048xf32> loc(#loc5262)
    %14956 = stablehlo.broadcast_in_dim %14938, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %14957 = stablehlo.broadcast_in_dim %14956, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %14958 = stablehlo.multiply %14957, %arg523 : tensor<1x2048xf32> loc(#loc5262)
    %14959 = stablehlo.add %14955, %14958 : tensor<1x2048xf32> loc(#loc5263)
    %14960 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14961 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14962 = stablehlo.broadcast_in_dim %14961, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14963 = stablehlo.add %14960, %14962 : tensor<1xf32> loc(#loc5259)
    %14964 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14965 = stablehlo.broadcast_in_dim %14964, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14966 = stablehlo.subtract %14963, %14965 : tensor<1xf32> loc(#loc5260)
    %14967 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14968 = stablehlo.broadcast_in_dim %14967, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14969 = stablehlo.power %14968, %14966 : tensor<1xf32> loc(#loc5261)
    %14970 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14971 = stablehlo.broadcast_in_dim %14970, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14972 = stablehlo.subtract %14971, %14969 : tensor<1xf32> loc(#loc5260)
    %14973 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14974 = stablehlo.broadcast_in_dim %14973, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14975 = stablehlo.multiply %14974, %14972 : tensor<1xf32> loc(#loc5236)
    %14976 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %14977 = stablehlo.broadcast_in_dim %14976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14978 = stablehlo.power %14977, %14963 : tensor<1xf32> loc(#loc5261)
    %14979 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14980 = stablehlo.broadcast_in_dim %14979, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14981 = stablehlo.subtract %14980, %14978 : tensor<1xf32> loc(#loc5260)
    %14982 = stablehlo.divide %14975, %14981 : tensor<1xf32> loc(#loc5253)
    %14983 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %14984 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14985 = stablehlo.broadcast_in_dim %14984, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %14986 = stablehlo.add %14983, %14985 : tensor<1xf32> loc(#loc5259)
    %14987 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14988 = stablehlo.broadcast_in_dim %14987, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14989 = stablehlo.subtract %14986, %14988 : tensor<1xf32> loc(#loc5260)
    %14990 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14991 = stablehlo.broadcast_in_dim %14990, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %14992 = stablehlo.power %14991, %14989 : tensor<1xf32> loc(#loc5261)
    %14993 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %14994 = stablehlo.broadcast_in_dim %14993, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %14995 = stablehlo.subtract %14994, %14992 : tensor<1xf32> loc(#loc5260)
    %14996 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %14997 = stablehlo.broadcast_in_dim %14996, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %14998 = stablehlo.multiply %14997, %14995 : tensor<1xf32> loc(#loc5236)
    %14999 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15000 = stablehlo.broadcast_in_dim %14999, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15001 = stablehlo.power %15000, %14986 : tensor<1xf32> loc(#loc5261)
    %15002 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15003 = stablehlo.broadcast_in_dim %15002, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15004 = stablehlo.subtract %15003, %15001 : tensor<1xf32> loc(#loc5260)
    %15005 = stablehlo.divide %14998, %15004 : tensor<1xf32> loc(#loc5253)
    %15006 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15007 = stablehlo.broadcast_in_dim %15006, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15008 = stablehlo.subtract %15007, %14982 : tensor<1xf32> loc(#loc5260)
    %15009 = stablehlo.broadcast_in_dim %15008, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15010 = stablehlo.broadcast_in_dim %15009, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15011 = stablehlo.multiply %15010, %4872 : tensor<1x2048xf32> loc(#loc5262)
    %15012 = stablehlo.broadcast_in_dim %14982, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15013 = stablehlo.broadcast_in_dim %15012, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15014 = stablehlo.multiply %15013, %arg348 : tensor<1x2048xf32> loc(#loc5262)
    %15015 = stablehlo.add %15011, %15014 : tensor<1x2048xf32> loc(#loc5263)
    %15016 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15017 = stablehlo.broadcast_in_dim %15016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15018 = stablehlo.subtract %15017, %15005 : tensor<1xf32> loc(#loc5260)
    %15019 = stablehlo.multiply %4872, %4872 : tensor<1x2048xf32> loc(#loc5264)
    %15020 = stablehlo.broadcast_in_dim %15018, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15021 = stablehlo.broadcast_in_dim %15020, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15022 = stablehlo.multiply %15021, %15019 : tensor<1x2048xf32> loc(#loc5262)
    %15023 = stablehlo.broadcast_in_dim %15005, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15024 = stablehlo.broadcast_in_dim %15023, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15025 = stablehlo.multiply %15024, %arg524 : tensor<1x2048xf32> loc(#loc5262)
    %15026 = stablehlo.add %15022, %15025 : tensor<1x2048xf32> loc(#loc5263)
    %15027 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15028 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15029 = stablehlo.broadcast_in_dim %15028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15030 = stablehlo.add %15027, %15029 : tensor<1xf32> loc(#loc5259)
    %15031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15032 = stablehlo.broadcast_in_dim %15031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15033 = stablehlo.subtract %15030, %15032 : tensor<1xf32> loc(#loc5260)
    %15034 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15035 = stablehlo.broadcast_in_dim %15034, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15036 = stablehlo.power %15035, %15033 : tensor<1xf32> loc(#loc5261)
    %15037 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15038 = stablehlo.broadcast_in_dim %15037, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15039 = stablehlo.subtract %15038, %15036 : tensor<1xf32> loc(#loc5260)
    %15040 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15041 = stablehlo.broadcast_in_dim %15040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15042 = stablehlo.multiply %15041, %15039 : tensor<1xf32> loc(#loc5236)
    %15043 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15044 = stablehlo.broadcast_in_dim %15043, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15045 = stablehlo.power %15044, %15030 : tensor<1xf32> loc(#loc5261)
    %15046 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15047 = stablehlo.broadcast_in_dim %15046, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15048 = stablehlo.subtract %15047, %15045 : tensor<1xf32> loc(#loc5260)
    %15049 = stablehlo.divide %15042, %15048 : tensor<1xf32> loc(#loc5253)
    %15050 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15051 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15052 = stablehlo.broadcast_in_dim %15051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15053 = stablehlo.add %15050, %15052 : tensor<1xf32> loc(#loc5259)
    %15054 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15055 = stablehlo.broadcast_in_dim %15054, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15056 = stablehlo.subtract %15053, %15055 : tensor<1xf32> loc(#loc5260)
    %15057 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15058 = stablehlo.broadcast_in_dim %15057, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15059 = stablehlo.power %15058, %15056 : tensor<1xf32> loc(#loc5261)
    %15060 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15061 = stablehlo.broadcast_in_dim %15060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15062 = stablehlo.subtract %15061, %15059 : tensor<1xf32> loc(#loc5260)
    %15063 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15064 = stablehlo.broadcast_in_dim %15063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15065 = stablehlo.multiply %15064, %15062 : tensor<1xf32> loc(#loc5236)
    %15066 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15067 = stablehlo.broadcast_in_dim %15066, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15068 = stablehlo.power %15067, %15053 : tensor<1xf32> loc(#loc5261)
    %15069 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15070 = stablehlo.broadcast_in_dim %15069, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15071 = stablehlo.subtract %15070, %15068 : tensor<1xf32> loc(#loc5260)
    %15072 = stablehlo.divide %15065, %15071 : tensor<1xf32> loc(#loc5253)
    %15073 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15074 = stablehlo.broadcast_in_dim %15073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15075 = stablehlo.subtract %15074, %15049 : tensor<1xf32> loc(#loc5260)
    %15076 = stablehlo.broadcast_in_dim %15075, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15077 = stablehlo.broadcast_in_dim %15076, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15078 = stablehlo.multiply %15077, %4876 : tensor<1x2048xf32> loc(#loc5262)
    %15079 = stablehlo.broadcast_in_dim %15049, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15080 = stablehlo.broadcast_in_dim %15079, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15081 = stablehlo.multiply %15080, %arg349 : tensor<1x2048xf32> loc(#loc5262)
    %15082 = stablehlo.add %15078, %15081 : tensor<1x2048xf32> loc(#loc5263)
    %15083 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15084 = stablehlo.broadcast_in_dim %15083, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15085 = stablehlo.subtract %15084, %15072 : tensor<1xf32> loc(#loc5260)
    %15086 = stablehlo.multiply %4876, %4876 : tensor<1x2048xf32> loc(#loc5264)
    %15087 = stablehlo.broadcast_in_dim %15085, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15088 = stablehlo.broadcast_in_dim %15087, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15089 = stablehlo.multiply %15088, %15086 : tensor<1x2048xf32> loc(#loc5262)
    %15090 = stablehlo.broadcast_in_dim %15072, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15091 = stablehlo.broadcast_in_dim %15090, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15092 = stablehlo.multiply %15091, %arg525 : tensor<1x2048xf32> loc(#loc5262)
    %15093 = stablehlo.add %15089, %15092 : tensor<1x2048xf32> loc(#loc5263)
    %15094 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15095 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15096 = stablehlo.broadcast_in_dim %15095, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15097 = stablehlo.add %15094, %15096 : tensor<1xf32> loc(#loc5259)
    %15098 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15099 = stablehlo.broadcast_in_dim %15098, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15100 = stablehlo.subtract %15097, %15099 : tensor<1xf32> loc(#loc5260)
    %15101 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15102 = stablehlo.broadcast_in_dim %15101, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15103 = stablehlo.power %15102, %15100 : tensor<1xf32> loc(#loc5261)
    %15104 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15105 = stablehlo.broadcast_in_dim %15104, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15106 = stablehlo.subtract %15105, %15103 : tensor<1xf32> loc(#loc5260)
    %15107 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15108 = stablehlo.broadcast_in_dim %15107, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15109 = stablehlo.multiply %15108, %15106 : tensor<1xf32> loc(#loc5236)
    %15110 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15111 = stablehlo.broadcast_in_dim %15110, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15112 = stablehlo.power %15111, %15097 : tensor<1xf32> loc(#loc5261)
    %15113 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15114 = stablehlo.broadcast_in_dim %15113, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15115 = stablehlo.subtract %15114, %15112 : tensor<1xf32> loc(#loc5260)
    %15116 = stablehlo.divide %15109, %15115 : tensor<1xf32> loc(#loc5253)
    %15117 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15118 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15119 = stablehlo.broadcast_in_dim %15118, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15120 = stablehlo.add %15117, %15119 : tensor<1xf32> loc(#loc5259)
    %15121 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15122 = stablehlo.broadcast_in_dim %15121, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15123 = stablehlo.subtract %15120, %15122 : tensor<1xf32> loc(#loc5260)
    %15124 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15125 = stablehlo.broadcast_in_dim %15124, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15126 = stablehlo.power %15125, %15123 : tensor<1xf32> loc(#loc5261)
    %15127 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15128 = stablehlo.broadcast_in_dim %15127, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15129 = stablehlo.subtract %15128, %15126 : tensor<1xf32> loc(#loc5260)
    %15130 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15131 = stablehlo.broadcast_in_dim %15130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15132 = stablehlo.multiply %15131, %15129 : tensor<1xf32> loc(#loc5236)
    %15133 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15134 = stablehlo.broadcast_in_dim %15133, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15135 = stablehlo.power %15134, %15120 : tensor<1xf32> loc(#loc5261)
    %15136 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15137 = stablehlo.broadcast_in_dim %15136, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15138 = stablehlo.subtract %15137, %15135 : tensor<1xf32> loc(#loc5260)
    %15139 = stablehlo.divide %15132, %15138 : tensor<1xf32> loc(#loc5253)
    %15140 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15141 = stablehlo.broadcast_in_dim %15140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15142 = stablehlo.subtract %15141, %15116 : tensor<1xf32> loc(#loc5260)
    %15143 = stablehlo.broadcast_in_dim %15142, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15144 = stablehlo.broadcast_in_dim %15143, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15145 = stablehlo.multiply %15144, %4880 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15146 = stablehlo.broadcast_in_dim %15116, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15147 = stablehlo.broadcast_in_dim %15146, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15148 = stablehlo.multiply %15147, %arg350 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15149 = stablehlo.add %15145, %15148 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %15150 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15151 = stablehlo.broadcast_in_dim %15150, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15152 = stablehlo.subtract %15151, %15139 : tensor<1xf32> loc(#loc5260)
    %15153 = stablehlo.multiply %4880, %4880 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %15154 = stablehlo.broadcast_in_dim %15152, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15155 = stablehlo.broadcast_in_dim %15154, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15156 = stablehlo.multiply %15155, %15153 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15157 = stablehlo.broadcast_in_dim %15139, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15158 = stablehlo.broadcast_in_dim %15157, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15159 = stablehlo.multiply %15158, %arg526 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15160 = stablehlo.add %15156, %15159 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %15161 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15162 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15163 = stablehlo.broadcast_in_dim %15162, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15164 = stablehlo.add %15161, %15163 : tensor<1xf32> loc(#loc5259)
    %15165 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15166 = stablehlo.broadcast_in_dim %15165, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15167 = stablehlo.subtract %15164, %15166 : tensor<1xf32> loc(#loc5260)
    %15168 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15169 = stablehlo.broadcast_in_dim %15168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15170 = stablehlo.power %15169, %15167 : tensor<1xf32> loc(#loc5261)
    %15171 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15172 = stablehlo.broadcast_in_dim %15171, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15173 = stablehlo.subtract %15172, %15170 : tensor<1xf32> loc(#loc5260)
    %15174 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15175 = stablehlo.broadcast_in_dim %15174, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15176 = stablehlo.multiply %15175, %15173 : tensor<1xf32> loc(#loc5236)
    %15177 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15178 = stablehlo.broadcast_in_dim %15177, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15179 = stablehlo.power %15178, %15164 : tensor<1xf32> loc(#loc5261)
    %15180 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15181 = stablehlo.broadcast_in_dim %15180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15182 = stablehlo.subtract %15181, %15179 : tensor<1xf32> loc(#loc5260)
    %15183 = stablehlo.divide %15176, %15182 : tensor<1xf32> loc(#loc5253)
    %15184 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15185 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15186 = stablehlo.broadcast_in_dim %15185, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15187 = stablehlo.add %15184, %15186 : tensor<1xf32> loc(#loc5259)
    %15188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15189 = stablehlo.broadcast_in_dim %15188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15190 = stablehlo.subtract %15187, %15189 : tensor<1xf32> loc(#loc5260)
    %15191 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15192 = stablehlo.broadcast_in_dim %15191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15193 = stablehlo.power %15192, %15190 : tensor<1xf32> loc(#loc5261)
    %15194 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15195 = stablehlo.broadcast_in_dim %15194, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15196 = stablehlo.subtract %15195, %15193 : tensor<1xf32> loc(#loc5260)
    %15197 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15198 = stablehlo.broadcast_in_dim %15197, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15199 = stablehlo.multiply %15198, %15196 : tensor<1xf32> loc(#loc5236)
    %15200 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15201 = stablehlo.broadcast_in_dim %15200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15202 = stablehlo.power %15201, %15187 : tensor<1xf32> loc(#loc5261)
    %15203 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15204 = stablehlo.broadcast_in_dim %15203, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15205 = stablehlo.subtract %15204, %15202 : tensor<1xf32> loc(#loc5260)
    %15206 = stablehlo.divide %15199, %15205 : tensor<1xf32> loc(#loc5253)
    %15207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15208 = stablehlo.broadcast_in_dim %15207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15209 = stablehlo.subtract %15208, %15183 : tensor<1xf32> loc(#loc5260)
    %15210 = stablehlo.broadcast_in_dim %15209, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15211 = stablehlo.broadcast_in_dim %15210, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15212 = stablehlo.multiply %15211, %4884 : tensor<1x64xf32> loc(#loc5262)
    %15213 = stablehlo.broadcast_in_dim %15183, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15214 = stablehlo.broadcast_in_dim %15213, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15215 = stablehlo.multiply %15214, %arg351 : tensor<1x64xf32> loc(#loc5262)
    %15216 = stablehlo.add %15212, %15215 : tensor<1x64xf32> loc(#loc5263)
    %15217 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15218 = stablehlo.broadcast_in_dim %15217, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15219 = stablehlo.subtract %15218, %15206 : tensor<1xf32> loc(#loc5260)
    %15220 = stablehlo.multiply %4884, %4884 : tensor<1x64xf32> loc(#loc5264)
    %15221 = stablehlo.broadcast_in_dim %15219, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15222 = stablehlo.broadcast_in_dim %15221, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15223 = stablehlo.multiply %15222, %15220 : tensor<1x64xf32> loc(#loc5262)
    %15224 = stablehlo.broadcast_in_dim %15206, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15225 = stablehlo.broadcast_in_dim %15224, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15226 = stablehlo.multiply %15225, %arg527 : tensor<1x64xf32> loc(#loc5262)
    %15227 = stablehlo.add %15223, %15226 : tensor<1x64xf32> loc(#loc5263)
    %15228 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15229 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15230 = stablehlo.broadcast_in_dim %15229, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15231 = stablehlo.add %15228, %15230 : tensor<1xf32> loc(#loc5259)
    %15232 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15233 = stablehlo.broadcast_in_dim %15232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15234 = stablehlo.subtract %15231, %15233 : tensor<1xf32> loc(#loc5260)
    %15235 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15236 = stablehlo.broadcast_in_dim %15235, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15237 = stablehlo.power %15236, %15234 : tensor<1xf32> loc(#loc5261)
    %15238 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15239 = stablehlo.broadcast_in_dim %15238, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15240 = stablehlo.subtract %15239, %15237 : tensor<1xf32> loc(#loc5260)
    %15241 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15242 = stablehlo.broadcast_in_dim %15241, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15243 = stablehlo.multiply %15242, %15240 : tensor<1xf32> loc(#loc5236)
    %15244 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15245 = stablehlo.broadcast_in_dim %15244, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15246 = stablehlo.power %15245, %15231 : tensor<1xf32> loc(#loc5261)
    %15247 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15248 = stablehlo.broadcast_in_dim %15247, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15249 = stablehlo.subtract %15248, %15246 : tensor<1xf32> loc(#loc5260)
    %15250 = stablehlo.divide %15243, %15249 : tensor<1xf32> loc(#loc5253)
    %15251 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15252 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15253 = stablehlo.broadcast_in_dim %15252, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15254 = stablehlo.add %15251, %15253 : tensor<1xf32> loc(#loc5259)
    %15255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15256 = stablehlo.broadcast_in_dim %15255, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15257 = stablehlo.subtract %15254, %15256 : tensor<1xf32> loc(#loc5260)
    %15258 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15259 = stablehlo.broadcast_in_dim %15258, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15260 = stablehlo.power %15259, %15257 : tensor<1xf32> loc(#loc5261)
    %15261 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15262 = stablehlo.broadcast_in_dim %15261, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15263 = stablehlo.subtract %15262, %15260 : tensor<1xf32> loc(#loc5260)
    %15264 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15265 = stablehlo.broadcast_in_dim %15264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15266 = stablehlo.multiply %15265, %15263 : tensor<1xf32> loc(#loc5236)
    %15267 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15268 = stablehlo.broadcast_in_dim %15267, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15269 = stablehlo.power %15268, %15254 : tensor<1xf32> loc(#loc5261)
    %15270 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15271 = stablehlo.broadcast_in_dim %15270, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15272 = stablehlo.subtract %15271, %15269 : tensor<1xf32> loc(#loc5260)
    %15273 = stablehlo.divide %15266, %15272 : tensor<1xf32> loc(#loc5253)
    %15274 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15275 = stablehlo.broadcast_in_dim %15274, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15276 = stablehlo.subtract %15275, %15250 : tensor<1xf32> loc(#loc5260)
    %15277 = stablehlo.broadcast_in_dim %15276, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %15278 = stablehlo.broadcast_in_dim %15277, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %15279 = stablehlo.multiply %15278, %4888 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %15280 = stablehlo.broadcast_in_dim %15250, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %15281 = stablehlo.broadcast_in_dim %15280, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %15282 = stablehlo.multiply %15281, %arg352 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %15283 = stablehlo.add %15279, %15282 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %15284 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15285 = stablehlo.broadcast_in_dim %15284, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15286 = stablehlo.subtract %15285, %15273 : tensor<1xf32> loc(#loc5260)
    %15287 = stablehlo.multiply %4888, %4888 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %15288 = stablehlo.broadcast_in_dim %15286, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %15289 = stablehlo.broadcast_in_dim %15288, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %15290 = stablehlo.multiply %15289, %15287 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %15291 = stablehlo.broadcast_in_dim %15273, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %15292 = stablehlo.broadcast_in_dim %15291, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %15293 = stablehlo.multiply %15292, %arg528 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %15294 = stablehlo.add %15290, %15293 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %15295 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15296 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15297 = stablehlo.broadcast_in_dim %15296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15298 = stablehlo.add %15295, %15297 : tensor<1xf32> loc(#loc5259)
    %15299 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15300 = stablehlo.broadcast_in_dim %15299, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15301 = stablehlo.subtract %15298, %15300 : tensor<1xf32> loc(#loc5260)
    %15302 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15303 = stablehlo.broadcast_in_dim %15302, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15304 = stablehlo.power %15303, %15301 : tensor<1xf32> loc(#loc5261)
    %15305 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15306 = stablehlo.broadcast_in_dim %15305, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15307 = stablehlo.subtract %15306, %15304 : tensor<1xf32> loc(#loc5260)
    %15308 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15309 = stablehlo.broadcast_in_dim %15308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15310 = stablehlo.multiply %15309, %15307 : tensor<1xf32> loc(#loc5236)
    %15311 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15312 = stablehlo.broadcast_in_dim %15311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15313 = stablehlo.power %15312, %15298 : tensor<1xf32> loc(#loc5261)
    %15314 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15315 = stablehlo.broadcast_in_dim %15314, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15316 = stablehlo.subtract %15315, %15313 : tensor<1xf32> loc(#loc5260)
    %15317 = stablehlo.divide %15310, %15316 : tensor<1xf32> loc(#loc5253)
    %15318 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15319 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15320 = stablehlo.broadcast_in_dim %15319, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15321 = stablehlo.add %15318, %15320 : tensor<1xf32> loc(#loc5259)
    %15322 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15323 = stablehlo.broadcast_in_dim %15322, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15324 = stablehlo.subtract %15321, %15323 : tensor<1xf32> loc(#loc5260)
    %15325 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15326 = stablehlo.broadcast_in_dim %15325, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15327 = stablehlo.power %15326, %15324 : tensor<1xf32> loc(#loc5261)
    %15328 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15329 = stablehlo.broadcast_in_dim %15328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15330 = stablehlo.subtract %15329, %15327 : tensor<1xf32> loc(#loc5260)
    %15331 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15332 = stablehlo.broadcast_in_dim %15331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15333 = stablehlo.multiply %15332, %15330 : tensor<1xf32> loc(#loc5236)
    %15334 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15335 = stablehlo.broadcast_in_dim %15334, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15336 = stablehlo.power %15335, %15321 : tensor<1xf32> loc(#loc5261)
    %15337 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15338 = stablehlo.broadcast_in_dim %15337, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15339 = stablehlo.subtract %15338, %15336 : tensor<1xf32> loc(#loc5260)
    %15340 = stablehlo.divide %15333, %15339 : tensor<1xf32> loc(#loc5253)
    %15341 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15342 = stablehlo.broadcast_in_dim %15341, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15343 = stablehlo.subtract %15342, %15317 : tensor<1xf32> loc(#loc5260)
    %15344 = stablehlo.broadcast_in_dim %15343, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15345 = stablehlo.broadcast_in_dim %15344, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %15346 = stablehlo.multiply %15345, %4892 : tensor<1x8192xf32> loc(#loc5262)
    %15347 = stablehlo.broadcast_in_dim %15317, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15348 = stablehlo.broadcast_in_dim %15347, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %15349 = stablehlo.multiply %15348, %arg353 : tensor<1x8192xf32> loc(#loc5262)
    %15350 = stablehlo.add %15346, %15349 : tensor<1x8192xf32> loc(#loc5263)
    %15351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15352 = stablehlo.broadcast_in_dim %15351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15353 = stablehlo.subtract %15352, %15340 : tensor<1xf32> loc(#loc5260)
    %15354 = stablehlo.multiply %4892, %4892 : tensor<1x8192xf32> loc(#loc5264)
    %15355 = stablehlo.broadcast_in_dim %15353, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15356 = stablehlo.broadcast_in_dim %15355, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %15357 = stablehlo.multiply %15356, %15354 : tensor<1x8192xf32> loc(#loc5262)
    %15358 = stablehlo.broadcast_in_dim %15340, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15359 = stablehlo.broadcast_in_dim %15358, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %15360 = stablehlo.multiply %15359, %arg529 : tensor<1x8192xf32> loc(#loc5262)
    %15361 = stablehlo.add %15357, %15360 : tensor<1x8192xf32> loc(#loc5263)
    %15362 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15363 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15364 = stablehlo.broadcast_in_dim %15363, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15365 = stablehlo.add %15362, %15364 : tensor<1xf32> loc(#loc5259)
    %15366 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15367 = stablehlo.broadcast_in_dim %15366, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15368 = stablehlo.subtract %15365, %15367 : tensor<1xf32> loc(#loc5260)
    %15369 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15370 = stablehlo.broadcast_in_dim %15369, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15371 = stablehlo.power %15370, %15368 : tensor<1xf32> loc(#loc5261)
    %15372 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15373 = stablehlo.broadcast_in_dim %15372, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15374 = stablehlo.subtract %15373, %15371 : tensor<1xf32> loc(#loc5260)
    %15375 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15376 = stablehlo.broadcast_in_dim %15375, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15377 = stablehlo.multiply %15376, %15374 : tensor<1xf32> loc(#loc5236)
    %15378 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15379 = stablehlo.broadcast_in_dim %15378, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15380 = stablehlo.power %15379, %15365 : tensor<1xf32> loc(#loc5261)
    %15381 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15382 = stablehlo.broadcast_in_dim %15381, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15383 = stablehlo.subtract %15382, %15380 : tensor<1xf32> loc(#loc5260)
    %15384 = stablehlo.divide %15377, %15383 : tensor<1xf32> loc(#loc5253)
    %15385 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15386 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15387 = stablehlo.broadcast_in_dim %15386, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15388 = stablehlo.add %15385, %15387 : tensor<1xf32> loc(#loc5259)
    %15389 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15390 = stablehlo.broadcast_in_dim %15389, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15391 = stablehlo.subtract %15388, %15390 : tensor<1xf32> loc(#loc5260)
    %15392 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15393 = stablehlo.broadcast_in_dim %15392, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15394 = stablehlo.power %15393, %15391 : tensor<1xf32> loc(#loc5261)
    %15395 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15396 = stablehlo.broadcast_in_dim %15395, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15397 = stablehlo.subtract %15396, %15394 : tensor<1xf32> loc(#loc5260)
    %15398 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15399 = stablehlo.broadcast_in_dim %15398, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15400 = stablehlo.multiply %15399, %15397 : tensor<1xf32> loc(#loc5236)
    %15401 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15402 = stablehlo.broadcast_in_dim %15401, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15403 = stablehlo.power %15402, %15388 : tensor<1xf32> loc(#loc5261)
    %15404 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15405 = stablehlo.broadcast_in_dim %15404, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15406 = stablehlo.subtract %15405, %15403 : tensor<1xf32> loc(#loc5260)
    %15407 = stablehlo.divide %15400, %15406 : tensor<1xf32> loc(#loc5253)
    %15408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15409 = stablehlo.broadcast_in_dim %15408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15410 = stablehlo.subtract %15409, %15384 : tensor<1xf32> loc(#loc5260)
    %15411 = stablehlo.broadcast_in_dim %15410, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15412 = stablehlo.broadcast_in_dim %15411, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %15413 = stablehlo.multiply %15412, %4896 : tensor<1x2048x8192xf32> loc(#loc5262)
    %15414 = stablehlo.broadcast_in_dim %15384, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15415 = stablehlo.broadcast_in_dim %15414, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %15416 = stablehlo.multiply %15415, %arg354 : tensor<1x2048x8192xf32> loc(#loc5262)
    %15417 = stablehlo.add %15413, %15416 : tensor<1x2048x8192xf32> loc(#loc5263)
    %15418 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15419 = stablehlo.broadcast_in_dim %15418, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15420 = stablehlo.subtract %15419, %15407 : tensor<1xf32> loc(#loc5260)
    %15421 = stablehlo.multiply %4896, %4896 : tensor<1x2048x8192xf32> loc(#loc5264)
    %15422 = stablehlo.broadcast_in_dim %15420, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15423 = stablehlo.broadcast_in_dim %15422, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %15424 = stablehlo.multiply %15423, %15421 : tensor<1x2048x8192xf32> loc(#loc5262)
    %15425 = stablehlo.broadcast_in_dim %15407, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15426 = stablehlo.broadcast_in_dim %15425, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %15427 = stablehlo.multiply %15426, %arg530 : tensor<1x2048x8192xf32> loc(#loc5262)
    %15428 = stablehlo.add %15424, %15427 : tensor<1x2048x8192xf32> loc(#loc5263)
    %15429 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15430 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15431 = stablehlo.broadcast_in_dim %15430, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15432 = stablehlo.add %15429, %15431 : tensor<1xf32> loc(#loc5259)
    %15433 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15434 = stablehlo.broadcast_in_dim %15433, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15435 = stablehlo.subtract %15432, %15434 : tensor<1xf32> loc(#loc5260)
    %15436 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15437 = stablehlo.broadcast_in_dim %15436, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15438 = stablehlo.power %15437, %15435 : tensor<1xf32> loc(#loc5261)
    %15439 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15440 = stablehlo.broadcast_in_dim %15439, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15441 = stablehlo.subtract %15440, %15438 : tensor<1xf32> loc(#loc5260)
    %15442 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15443 = stablehlo.broadcast_in_dim %15442, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15444 = stablehlo.multiply %15443, %15441 : tensor<1xf32> loc(#loc5236)
    %15445 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15446 = stablehlo.broadcast_in_dim %15445, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15447 = stablehlo.power %15446, %15432 : tensor<1xf32> loc(#loc5261)
    %15448 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15449 = stablehlo.broadcast_in_dim %15448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15450 = stablehlo.subtract %15449, %15447 : tensor<1xf32> loc(#loc5260)
    %15451 = stablehlo.divide %15444, %15450 : tensor<1xf32> loc(#loc5253)
    %15452 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15453 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15454 = stablehlo.broadcast_in_dim %15453, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15455 = stablehlo.add %15452, %15454 : tensor<1xf32> loc(#loc5259)
    %15456 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15457 = stablehlo.broadcast_in_dim %15456, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15458 = stablehlo.subtract %15455, %15457 : tensor<1xf32> loc(#loc5260)
    %15459 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15460 = stablehlo.broadcast_in_dim %15459, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15461 = stablehlo.power %15460, %15458 : tensor<1xf32> loc(#loc5261)
    %15462 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15463 = stablehlo.broadcast_in_dim %15462, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15464 = stablehlo.subtract %15463, %15461 : tensor<1xf32> loc(#loc5260)
    %15465 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15466 = stablehlo.broadcast_in_dim %15465, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15467 = stablehlo.multiply %15466, %15464 : tensor<1xf32> loc(#loc5236)
    %15468 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15469 = stablehlo.broadcast_in_dim %15468, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15470 = stablehlo.power %15469, %15455 : tensor<1xf32> loc(#loc5261)
    %15471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15472 = stablehlo.broadcast_in_dim %15471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15473 = stablehlo.subtract %15472, %15470 : tensor<1xf32> loc(#loc5260)
    %15474 = stablehlo.divide %15467, %15473 : tensor<1xf32> loc(#loc5253)
    %15475 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15476 = stablehlo.broadcast_in_dim %15475, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15477 = stablehlo.subtract %15476, %15451 : tensor<1xf32> loc(#loc5260)
    %15478 = stablehlo.broadcast_in_dim %15477, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15479 = stablehlo.broadcast_in_dim %15478, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15480 = stablehlo.multiply %15479, %4900 : tensor<1x2048xf32> loc(#loc5262)
    %15481 = stablehlo.broadcast_in_dim %15451, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15482 = stablehlo.broadcast_in_dim %15481, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15483 = stablehlo.multiply %15482, %arg355 : tensor<1x2048xf32> loc(#loc5262)
    %15484 = stablehlo.add %15480, %15483 : tensor<1x2048xf32> loc(#loc5263)
    %15485 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15486 = stablehlo.broadcast_in_dim %15485, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15487 = stablehlo.subtract %15486, %15474 : tensor<1xf32> loc(#loc5260)
    %15488 = stablehlo.multiply %4900, %4900 : tensor<1x2048xf32> loc(#loc5264)
    %15489 = stablehlo.broadcast_in_dim %15487, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15490 = stablehlo.broadcast_in_dim %15489, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15491 = stablehlo.multiply %15490, %15488 : tensor<1x2048xf32> loc(#loc5262)
    %15492 = stablehlo.broadcast_in_dim %15474, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15493 = stablehlo.broadcast_in_dim %15492, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15494 = stablehlo.multiply %15493, %arg531 : tensor<1x2048xf32> loc(#loc5262)
    %15495 = stablehlo.add %15491, %15494 : tensor<1x2048xf32> loc(#loc5263)
    %15496 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15497 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15498 = stablehlo.broadcast_in_dim %15497, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15499 = stablehlo.add %15496, %15498 : tensor<1xf32> loc(#loc5259)
    %15500 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15501 = stablehlo.broadcast_in_dim %15500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15502 = stablehlo.subtract %15499, %15501 : tensor<1xf32> loc(#loc5260)
    %15503 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15504 = stablehlo.broadcast_in_dim %15503, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15505 = stablehlo.power %15504, %15502 : tensor<1xf32> loc(#loc5261)
    %15506 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15507 = stablehlo.broadcast_in_dim %15506, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15508 = stablehlo.subtract %15507, %15505 : tensor<1xf32> loc(#loc5260)
    %15509 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15510 = stablehlo.broadcast_in_dim %15509, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15511 = stablehlo.multiply %15510, %15508 : tensor<1xf32> loc(#loc5236)
    %15512 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15513 = stablehlo.broadcast_in_dim %15512, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15514 = stablehlo.power %15513, %15499 : tensor<1xf32> loc(#loc5261)
    %15515 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15516 = stablehlo.broadcast_in_dim %15515, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15517 = stablehlo.subtract %15516, %15514 : tensor<1xf32> loc(#loc5260)
    %15518 = stablehlo.divide %15511, %15517 : tensor<1xf32> loc(#loc5253)
    %15519 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15520 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15521 = stablehlo.broadcast_in_dim %15520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15522 = stablehlo.add %15519, %15521 : tensor<1xf32> loc(#loc5259)
    %15523 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15524 = stablehlo.broadcast_in_dim %15523, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15525 = stablehlo.subtract %15522, %15524 : tensor<1xf32> loc(#loc5260)
    %15526 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15527 = stablehlo.broadcast_in_dim %15526, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15528 = stablehlo.power %15527, %15525 : tensor<1xf32> loc(#loc5261)
    %15529 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15530 = stablehlo.broadcast_in_dim %15529, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15531 = stablehlo.subtract %15530, %15528 : tensor<1xf32> loc(#loc5260)
    %15532 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15533 = stablehlo.broadcast_in_dim %15532, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15534 = stablehlo.multiply %15533, %15531 : tensor<1xf32> loc(#loc5236)
    %15535 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15536 = stablehlo.broadcast_in_dim %15535, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15537 = stablehlo.power %15536, %15522 : tensor<1xf32> loc(#loc5261)
    %15538 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15539 = stablehlo.broadcast_in_dim %15538, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15540 = stablehlo.subtract %15539, %15537 : tensor<1xf32> loc(#loc5260)
    %15541 = stablehlo.divide %15534, %15540 : tensor<1xf32> loc(#loc5253)
    %15542 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15543 = stablehlo.broadcast_in_dim %15542, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15544 = stablehlo.subtract %15543, %15518 : tensor<1xf32> loc(#loc5260)
    %15545 = stablehlo.broadcast_in_dim %15544, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15546 = stablehlo.broadcast_in_dim %15545, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %15547 = stablehlo.multiply %15546, %4904 : tensor<1x8192x2048xf32> loc(#loc5262)
    %15548 = stablehlo.broadcast_in_dim %15518, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15549 = stablehlo.broadcast_in_dim %15548, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %15550 = stablehlo.multiply %15549, %arg356 : tensor<1x8192x2048xf32> loc(#loc5262)
    %15551 = stablehlo.add %15547, %15550 : tensor<1x8192x2048xf32> loc(#loc5263)
    %15552 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15553 = stablehlo.broadcast_in_dim %15552, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15554 = stablehlo.subtract %15553, %15541 : tensor<1xf32> loc(#loc5260)
    %15555 = stablehlo.multiply %4904, %4904 : tensor<1x8192x2048xf32> loc(#loc5264)
    %15556 = stablehlo.broadcast_in_dim %15554, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15557 = stablehlo.broadcast_in_dim %15556, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %15558 = stablehlo.multiply %15557, %15555 : tensor<1x8192x2048xf32> loc(#loc5262)
    %15559 = stablehlo.broadcast_in_dim %15541, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %15560 = stablehlo.broadcast_in_dim %15559, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %15561 = stablehlo.multiply %15560, %arg532 : tensor<1x8192x2048xf32> loc(#loc5262)
    %15562 = stablehlo.add %15558, %15561 : tensor<1x8192x2048xf32> loc(#loc5263)
    %15563 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15564 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15565 = stablehlo.broadcast_in_dim %15564, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15566 = stablehlo.add %15563, %15565 : tensor<1xf32> loc(#loc5259)
    %15567 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15568 = stablehlo.broadcast_in_dim %15567, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15569 = stablehlo.subtract %15566, %15568 : tensor<1xf32> loc(#loc5260)
    %15570 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15571 = stablehlo.broadcast_in_dim %15570, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15572 = stablehlo.power %15571, %15569 : tensor<1xf32> loc(#loc5261)
    %15573 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15574 = stablehlo.broadcast_in_dim %15573, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15575 = stablehlo.subtract %15574, %15572 : tensor<1xf32> loc(#loc5260)
    %15576 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15577 = stablehlo.broadcast_in_dim %15576, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15578 = stablehlo.multiply %15577, %15575 : tensor<1xf32> loc(#loc5236)
    %15579 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15580 = stablehlo.broadcast_in_dim %15579, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15581 = stablehlo.power %15580, %15566 : tensor<1xf32> loc(#loc5261)
    %15582 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15583 = stablehlo.broadcast_in_dim %15582, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15584 = stablehlo.subtract %15583, %15581 : tensor<1xf32> loc(#loc5260)
    %15585 = stablehlo.divide %15578, %15584 : tensor<1xf32> loc(#loc5253)
    %15586 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15587 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15588 = stablehlo.broadcast_in_dim %15587, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15589 = stablehlo.add %15586, %15588 : tensor<1xf32> loc(#loc5259)
    %15590 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15591 = stablehlo.broadcast_in_dim %15590, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15592 = stablehlo.subtract %15589, %15591 : tensor<1xf32> loc(#loc5260)
    %15593 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15594 = stablehlo.broadcast_in_dim %15593, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15595 = stablehlo.power %15594, %15592 : tensor<1xf32> loc(#loc5261)
    %15596 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15597 = stablehlo.broadcast_in_dim %15596, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15598 = stablehlo.subtract %15597, %15595 : tensor<1xf32> loc(#loc5260)
    %15599 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15600 = stablehlo.broadcast_in_dim %15599, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15601 = stablehlo.multiply %15600, %15598 : tensor<1xf32> loc(#loc5236)
    %15602 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15603 = stablehlo.broadcast_in_dim %15602, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15604 = stablehlo.power %15603, %15589 : tensor<1xf32> loc(#loc5261)
    %15605 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15606 = stablehlo.broadcast_in_dim %15605, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15607 = stablehlo.subtract %15606, %15604 : tensor<1xf32> loc(#loc5260)
    %15608 = stablehlo.divide %15601, %15607 : tensor<1xf32> loc(#loc5253)
    %15609 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15610 = stablehlo.broadcast_in_dim %15609, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15611 = stablehlo.subtract %15610, %15585 : tensor<1xf32> loc(#loc5260)
    %15612 = stablehlo.broadcast_in_dim %15611, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15613 = stablehlo.broadcast_in_dim %15612, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15614 = stablehlo.multiply %15613, %4908 : tensor<1x2048xf32> loc(#loc5262)
    %15615 = stablehlo.broadcast_in_dim %15585, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15616 = stablehlo.broadcast_in_dim %15615, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15617 = stablehlo.multiply %15616, %arg357 : tensor<1x2048xf32> loc(#loc5262)
    %15618 = stablehlo.add %15614, %15617 : tensor<1x2048xf32> loc(#loc5263)
    %15619 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15620 = stablehlo.broadcast_in_dim %15619, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15621 = stablehlo.subtract %15620, %15608 : tensor<1xf32> loc(#loc5260)
    %15622 = stablehlo.multiply %4908, %4908 : tensor<1x2048xf32> loc(#loc5264)
    %15623 = stablehlo.broadcast_in_dim %15621, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15624 = stablehlo.broadcast_in_dim %15623, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15625 = stablehlo.multiply %15624, %15622 : tensor<1x2048xf32> loc(#loc5262)
    %15626 = stablehlo.broadcast_in_dim %15608, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15627 = stablehlo.broadcast_in_dim %15626, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15628 = stablehlo.multiply %15627, %arg533 : tensor<1x2048xf32> loc(#loc5262)
    %15629 = stablehlo.add %15625, %15628 : tensor<1x2048xf32> loc(#loc5263)
    %15630 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15632 = stablehlo.broadcast_in_dim %15631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15633 = stablehlo.add %15630, %15632 : tensor<1xf32> loc(#loc5259)
    %15634 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15635 = stablehlo.broadcast_in_dim %15634, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15636 = stablehlo.subtract %15633, %15635 : tensor<1xf32> loc(#loc5260)
    %15637 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15638 = stablehlo.broadcast_in_dim %15637, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15639 = stablehlo.power %15638, %15636 : tensor<1xf32> loc(#loc5261)
    %15640 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15641 = stablehlo.broadcast_in_dim %15640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15642 = stablehlo.subtract %15641, %15639 : tensor<1xf32> loc(#loc5260)
    %15643 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15644 = stablehlo.broadcast_in_dim %15643, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15645 = stablehlo.multiply %15644, %15642 : tensor<1xf32> loc(#loc5236)
    %15646 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15647 = stablehlo.broadcast_in_dim %15646, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15648 = stablehlo.power %15647, %15633 : tensor<1xf32> loc(#loc5261)
    %15649 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15650 = stablehlo.broadcast_in_dim %15649, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15651 = stablehlo.subtract %15650, %15648 : tensor<1xf32> loc(#loc5260)
    %15652 = stablehlo.divide %15645, %15651 : tensor<1xf32> loc(#loc5253)
    %15653 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15654 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15655 = stablehlo.broadcast_in_dim %15654, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15656 = stablehlo.add %15653, %15655 : tensor<1xf32> loc(#loc5259)
    %15657 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15658 = stablehlo.broadcast_in_dim %15657, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15659 = stablehlo.subtract %15656, %15658 : tensor<1xf32> loc(#loc5260)
    %15660 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15661 = stablehlo.broadcast_in_dim %15660, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15662 = stablehlo.power %15661, %15659 : tensor<1xf32> loc(#loc5261)
    %15663 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15664 = stablehlo.broadcast_in_dim %15663, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15665 = stablehlo.subtract %15664, %15662 : tensor<1xf32> loc(#loc5260)
    %15666 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15667 = stablehlo.broadcast_in_dim %15666, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15668 = stablehlo.multiply %15667, %15665 : tensor<1xf32> loc(#loc5236)
    %15669 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15670 = stablehlo.broadcast_in_dim %15669, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15671 = stablehlo.power %15670, %15656 : tensor<1xf32> loc(#loc5261)
    %15672 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15673 = stablehlo.broadcast_in_dim %15672, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15674 = stablehlo.subtract %15673, %15671 : tensor<1xf32> loc(#loc5260)
    %15675 = stablehlo.divide %15668, %15674 : tensor<1xf32> loc(#loc5253)
    %15676 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15677 = stablehlo.broadcast_in_dim %15676, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15678 = stablehlo.subtract %15677, %15652 : tensor<1xf32> loc(#loc5260)
    %15679 = stablehlo.broadcast_in_dim %15678, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15680 = stablehlo.broadcast_in_dim %15679, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15681 = stablehlo.multiply %15680, %4912 : tensor<1x2048xf32> loc(#loc5262)
    %15682 = stablehlo.broadcast_in_dim %15652, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15683 = stablehlo.broadcast_in_dim %15682, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15684 = stablehlo.multiply %15683, %arg358 : tensor<1x2048xf32> loc(#loc5262)
    %15685 = stablehlo.add %15681, %15684 : tensor<1x2048xf32> loc(#loc5263)
    %15686 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15687 = stablehlo.broadcast_in_dim %15686, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15688 = stablehlo.subtract %15687, %15675 : tensor<1xf32> loc(#loc5260)
    %15689 = stablehlo.multiply %4912, %4912 : tensor<1x2048xf32> loc(#loc5264)
    %15690 = stablehlo.broadcast_in_dim %15688, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15691 = stablehlo.broadcast_in_dim %15690, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15692 = stablehlo.multiply %15691, %15689 : tensor<1x2048xf32> loc(#loc5262)
    %15693 = stablehlo.broadcast_in_dim %15675, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15694 = stablehlo.broadcast_in_dim %15693, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15695 = stablehlo.multiply %15694, %arg534 : tensor<1x2048xf32> loc(#loc5262)
    %15696 = stablehlo.add %15692, %15695 : tensor<1x2048xf32> loc(#loc5263)
    %15697 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15698 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15699 = stablehlo.broadcast_in_dim %15698, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15700 = stablehlo.add %15697, %15699 : tensor<1xf32> loc(#loc5259)
    %15701 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15702 = stablehlo.broadcast_in_dim %15701, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15703 = stablehlo.subtract %15700, %15702 : tensor<1xf32> loc(#loc5260)
    %15704 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15705 = stablehlo.broadcast_in_dim %15704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15706 = stablehlo.power %15705, %15703 : tensor<1xf32> loc(#loc5261)
    %15707 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15708 = stablehlo.broadcast_in_dim %15707, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15709 = stablehlo.subtract %15708, %15706 : tensor<1xf32> loc(#loc5260)
    %15710 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15711 = stablehlo.broadcast_in_dim %15710, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15712 = stablehlo.multiply %15711, %15709 : tensor<1xf32> loc(#loc5236)
    %15713 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15714 = stablehlo.broadcast_in_dim %15713, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15715 = stablehlo.power %15714, %15700 : tensor<1xf32> loc(#loc5261)
    %15716 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15717 = stablehlo.broadcast_in_dim %15716, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15718 = stablehlo.subtract %15717, %15715 : tensor<1xf32> loc(#loc5260)
    %15719 = stablehlo.divide %15712, %15718 : tensor<1xf32> loc(#loc5253)
    %15720 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15721 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15722 = stablehlo.broadcast_in_dim %15721, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15723 = stablehlo.add %15720, %15722 : tensor<1xf32> loc(#loc5259)
    %15724 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15725 = stablehlo.broadcast_in_dim %15724, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15726 = stablehlo.subtract %15723, %15725 : tensor<1xf32> loc(#loc5260)
    %15727 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15728 = stablehlo.broadcast_in_dim %15727, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15729 = stablehlo.power %15728, %15726 : tensor<1xf32> loc(#loc5261)
    %15730 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15731 = stablehlo.broadcast_in_dim %15730, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15732 = stablehlo.subtract %15731, %15729 : tensor<1xf32> loc(#loc5260)
    %15733 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15734 = stablehlo.broadcast_in_dim %15733, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15735 = stablehlo.multiply %15734, %15732 : tensor<1xf32> loc(#loc5236)
    %15736 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15737 = stablehlo.broadcast_in_dim %15736, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15738 = stablehlo.power %15737, %15723 : tensor<1xf32> loc(#loc5261)
    %15739 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15740 = stablehlo.broadcast_in_dim %15739, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15741 = stablehlo.subtract %15740, %15738 : tensor<1xf32> loc(#loc5260)
    %15742 = stablehlo.divide %15735, %15741 : tensor<1xf32> loc(#loc5253)
    %15743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15744 = stablehlo.broadcast_in_dim %15743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15745 = stablehlo.subtract %15744, %15719 : tensor<1xf32> loc(#loc5260)
    %15746 = stablehlo.broadcast_in_dim %15745, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15747 = stablehlo.broadcast_in_dim %15746, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15748 = stablehlo.multiply %15747, %4916 : tensor<1x2048xf32> loc(#loc5262)
    %15749 = stablehlo.broadcast_in_dim %15719, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15750 = stablehlo.broadcast_in_dim %15749, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15751 = stablehlo.multiply %15750, %arg359 : tensor<1x2048xf32> loc(#loc5262)
    %15752 = stablehlo.add %15748, %15751 : tensor<1x2048xf32> loc(#loc5263)
    %15753 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15754 = stablehlo.broadcast_in_dim %15753, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15755 = stablehlo.subtract %15754, %15742 : tensor<1xf32> loc(#loc5260)
    %15756 = stablehlo.multiply %4916, %4916 : tensor<1x2048xf32> loc(#loc5264)
    %15757 = stablehlo.broadcast_in_dim %15755, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15758 = stablehlo.broadcast_in_dim %15757, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15759 = stablehlo.multiply %15758, %15756 : tensor<1x2048xf32> loc(#loc5262)
    %15760 = stablehlo.broadcast_in_dim %15742, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15761 = stablehlo.broadcast_in_dim %15760, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15762 = stablehlo.multiply %15761, %arg535 : tensor<1x2048xf32> loc(#loc5262)
    %15763 = stablehlo.add %15759, %15762 : tensor<1x2048xf32> loc(#loc5263)
    %15764 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15765 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15766 = stablehlo.broadcast_in_dim %15765, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15767 = stablehlo.add %15764, %15766 : tensor<1xf32> loc(#loc5259)
    %15768 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15769 = stablehlo.broadcast_in_dim %15768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15770 = stablehlo.subtract %15767, %15769 : tensor<1xf32> loc(#loc5260)
    %15771 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15772 = stablehlo.broadcast_in_dim %15771, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15773 = stablehlo.power %15772, %15770 : tensor<1xf32> loc(#loc5261)
    %15774 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15775 = stablehlo.broadcast_in_dim %15774, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15776 = stablehlo.subtract %15775, %15773 : tensor<1xf32> loc(#loc5260)
    %15777 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15778 = stablehlo.broadcast_in_dim %15777, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15779 = stablehlo.multiply %15778, %15776 : tensor<1xf32> loc(#loc5236)
    %15780 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15781 = stablehlo.broadcast_in_dim %15780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15782 = stablehlo.power %15781, %15767 : tensor<1xf32> loc(#loc5261)
    %15783 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15784 = stablehlo.broadcast_in_dim %15783, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15785 = stablehlo.subtract %15784, %15782 : tensor<1xf32> loc(#loc5260)
    %15786 = stablehlo.divide %15779, %15785 : tensor<1xf32> loc(#loc5253)
    %15787 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15788 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15789 = stablehlo.broadcast_in_dim %15788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15790 = stablehlo.add %15787, %15789 : tensor<1xf32> loc(#loc5259)
    %15791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15792 = stablehlo.broadcast_in_dim %15791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15793 = stablehlo.subtract %15790, %15792 : tensor<1xf32> loc(#loc5260)
    %15794 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15795 = stablehlo.broadcast_in_dim %15794, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15796 = stablehlo.power %15795, %15793 : tensor<1xf32> loc(#loc5261)
    %15797 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15798 = stablehlo.broadcast_in_dim %15797, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15799 = stablehlo.subtract %15798, %15796 : tensor<1xf32> loc(#loc5260)
    %15800 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15801 = stablehlo.broadcast_in_dim %15800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15802 = stablehlo.multiply %15801, %15799 : tensor<1xf32> loc(#loc5236)
    %15803 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15804 = stablehlo.broadcast_in_dim %15803, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15805 = stablehlo.power %15804, %15790 : tensor<1xf32> loc(#loc5261)
    %15806 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15807 = stablehlo.broadcast_in_dim %15806, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15808 = stablehlo.subtract %15807, %15805 : tensor<1xf32> loc(#loc5260)
    %15809 = stablehlo.divide %15802, %15808 : tensor<1xf32> loc(#loc5253)
    %15810 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15811 = stablehlo.broadcast_in_dim %15810, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15812 = stablehlo.subtract %15811, %15786 : tensor<1xf32> loc(#loc5260)
    %15813 = stablehlo.broadcast_in_dim %15812, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15814 = stablehlo.broadcast_in_dim %15813, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15815 = stablehlo.multiply %15814, %4920 : tensor<1x2048xf32> loc(#loc5262)
    %15816 = stablehlo.broadcast_in_dim %15786, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15817 = stablehlo.broadcast_in_dim %15816, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15818 = stablehlo.multiply %15817, %arg360 : tensor<1x2048xf32> loc(#loc5262)
    %15819 = stablehlo.add %15815, %15818 : tensor<1x2048xf32> loc(#loc5263)
    %15820 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15821 = stablehlo.broadcast_in_dim %15820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15822 = stablehlo.subtract %15821, %15809 : tensor<1xf32> loc(#loc5260)
    %15823 = stablehlo.multiply %4920, %4920 : tensor<1x2048xf32> loc(#loc5264)
    %15824 = stablehlo.broadcast_in_dim %15822, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15825 = stablehlo.broadcast_in_dim %15824, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15826 = stablehlo.multiply %15825, %15823 : tensor<1x2048xf32> loc(#loc5262)
    %15827 = stablehlo.broadcast_in_dim %15809, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15828 = stablehlo.broadcast_in_dim %15827, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %15829 = stablehlo.multiply %15828, %arg536 : tensor<1x2048xf32> loc(#loc5262)
    %15830 = stablehlo.add %15826, %15829 : tensor<1x2048xf32> loc(#loc5263)
    %15831 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15832 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15833 = stablehlo.broadcast_in_dim %15832, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15834 = stablehlo.add %15831, %15833 : tensor<1xf32> loc(#loc5259)
    %15835 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15836 = stablehlo.broadcast_in_dim %15835, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15837 = stablehlo.subtract %15834, %15836 : tensor<1xf32> loc(#loc5260)
    %15838 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15839 = stablehlo.broadcast_in_dim %15838, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15840 = stablehlo.power %15839, %15837 : tensor<1xf32> loc(#loc5261)
    %15841 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15842 = stablehlo.broadcast_in_dim %15841, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15843 = stablehlo.subtract %15842, %15840 : tensor<1xf32> loc(#loc5260)
    %15844 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15845 = stablehlo.broadcast_in_dim %15844, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15846 = stablehlo.multiply %15845, %15843 : tensor<1xf32> loc(#loc5236)
    %15847 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15848 = stablehlo.broadcast_in_dim %15847, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15849 = stablehlo.power %15848, %15834 : tensor<1xf32> loc(#loc5261)
    %15850 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15851 = stablehlo.broadcast_in_dim %15850, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15852 = stablehlo.subtract %15851, %15849 : tensor<1xf32> loc(#loc5260)
    %15853 = stablehlo.divide %15846, %15852 : tensor<1xf32> loc(#loc5253)
    %15854 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15855 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15856 = stablehlo.broadcast_in_dim %15855, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15857 = stablehlo.add %15854, %15856 : tensor<1xf32> loc(#loc5259)
    %15858 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15859 = stablehlo.broadcast_in_dim %15858, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15860 = stablehlo.subtract %15857, %15859 : tensor<1xf32> loc(#loc5260)
    %15861 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15862 = stablehlo.broadcast_in_dim %15861, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15863 = stablehlo.power %15862, %15860 : tensor<1xf32> loc(#loc5261)
    %15864 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15865 = stablehlo.broadcast_in_dim %15864, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15866 = stablehlo.subtract %15865, %15863 : tensor<1xf32> loc(#loc5260)
    %15867 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15868 = stablehlo.broadcast_in_dim %15867, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15869 = stablehlo.multiply %15868, %15866 : tensor<1xf32> loc(#loc5236)
    %15870 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15871 = stablehlo.broadcast_in_dim %15870, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15872 = stablehlo.power %15871, %15857 : tensor<1xf32> loc(#loc5261)
    %15873 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15874 = stablehlo.broadcast_in_dim %15873, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15875 = stablehlo.subtract %15874, %15872 : tensor<1xf32> loc(#loc5260)
    %15876 = stablehlo.divide %15869, %15875 : tensor<1xf32> loc(#loc5253)
    %15877 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15878 = stablehlo.broadcast_in_dim %15877, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15879 = stablehlo.subtract %15878, %15853 : tensor<1xf32> loc(#loc5260)
    %15880 = stablehlo.broadcast_in_dim %15879, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15881 = stablehlo.broadcast_in_dim %15880, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15882 = stablehlo.multiply %15881, %4924 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15883 = stablehlo.broadcast_in_dim %15853, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15884 = stablehlo.broadcast_in_dim %15883, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15885 = stablehlo.multiply %15884, %arg361 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15886 = stablehlo.add %15882, %15885 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %15887 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15888 = stablehlo.broadcast_in_dim %15887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15889 = stablehlo.subtract %15888, %15876 : tensor<1xf32> loc(#loc5260)
    %15890 = stablehlo.multiply %4924, %4924 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %15891 = stablehlo.broadcast_in_dim %15889, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15892 = stablehlo.broadcast_in_dim %15891, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15893 = stablehlo.multiply %15892, %15890 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15894 = stablehlo.broadcast_in_dim %15876, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %15895 = stablehlo.broadcast_in_dim %15894, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15896 = stablehlo.multiply %15895, %arg537 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %15897 = stablehlo.add %15893, %15896 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %15898 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15899 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15900 = stablehlo.broadcast_in_dim %15899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15901 = stablehlo.add %15898, %15900 : tensor<1xf32> loc(#loc5259)
    %15902 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15903 = stablehlo.broadcast_in_dim %15902, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15904 = stablehlo.subtract %15901, %15903 : tensor<1xf32> loc(#loc5260)
    %15905 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15906 = stablehlo.broadcast_in_dim %15905, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15907 = stablehlo.power %15906, %15904 : tensor<1xf32> loc(#loc5261)
    %15908 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15909 = stablehlo.broadcast_in_dim %15908, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15910 = stablehlo.subtract %15909, %15907 : tensor<1xf32> loc(#loc5260)
    %15911 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15912 = stablehlo.broadcast_in_dim %15911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15913 = stablehlo.multiply %15912, %15910 : tensor<1xf32> loc(#loc5236)
    %15914 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15915 = stablehlo.broadcast_in_dim %15914, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15916 = stablehlo.power %15915, %15901 : tensor<1xf32> loc(#loc5261)
    %15917 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15918 = stablehlo.broadcast_in_dim %15917, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15919 = stablehlo.subtract %15918, %15916 : tensor<1xf32> loc(#loc5260)
    %15920 = stablehlo.divide %15913, %15919 : tensor<1xf32> loc(#loc5253)
    %15921 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15922 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15923 = stablehlo.broadcast_in_dim %15922, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15924 = stablehlo.add %15921, %15923 : tensor<1xf32> loc(#loc5259)
    %15925 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15926 = stablehlo.broadcast_in_dim %15925, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15927 = stablehlo.subtract %15924, %15926 : tensor<1xf32> loc(#loc5260)
    %15928 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15929 = stablehlo.broadcast_in_dim %15928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15930 = stablehlo.power %15929, %15927 : tensor<1xf32> loc(#loc5261)
    %15931 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15932 = stablehlo.broadcast_in_dim %15931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15933 = stablehlo.subtract %15932, %15930 : tensor<1xf32> loc(#loc5260)
    %15934 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15935 = stablehlo.broadcast_in_dim %15934, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15936 = stablehlo.multiply %15935, %15933 : tensor<1xf32> loc(#loc5236)
    %15937 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15938 = stablehlo.broadcast_in_dim %15937, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15939 = stablehlo.power %15938, %15924 : tensor<1xf32> loc(#loc5261)
    %15940 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15941 = stablehlo.broadcast_in_dim %15940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15942 = stablehlo.subtract %15941, %15939 : tensor<1xf32> loc(#loc5260)
    %15943 = stablehlo.divide %15936, %15942 : tensor<1xf32> loc(#loc5253)
    %15944 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15945 = stablehlo.broadcast_in_dim %15944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15946 = stablehlo.subtract %15945, %15920 : tensor<1xf32> loc(#loc5260)
    %15947 = stablehlo.broadcast_in_dim %15946, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15948 = stablehlo.broadcast_in_dim %15947, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15949 = stablehlo.multiply %15948, %4928 : tensor<1x64xf32> loc(#loc5262)
    %15950 = stablehlo.broadcast_in_dim %15920, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15951 = stablehlo.broadcast_in_dim %15950, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15952 = stablehlo.multiply %15951, %arg362 : tensor<1x64xf32> loc(#loc5262)
    %15953 = stablehlo.add %15949, %15952 : tensor<1x64xf32> loc(#loc5263)
    %15954 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15955 = stablehlo.broadcast_in_dim %15954, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15956 = stablehlo.subtract %15955, %15943 : tensor<1xf32> loc(#loc5260)
    %15957 = stablehlo.multiply %4928, %4928 : tensor<1x64xf32> loc(#loc5264)
    %15958 = stablehlo.broadcast_in_dim %15956, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15959 = stablehlo.broadcast_in_dim %15958, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15960 = stablehlo.multiply %15959, %15957 : tensor<1x64xf32> loc(#loc5262)
    %15961 = stablehlo.broadcast_in_dim %15943, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %15962 = stablehlo.broadcast_in_dim %15961, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %15963 = stablehlo.multiply %15962, %arg538 : tensor<1x64xf32> loc(#loc5262)
    %15964 = stablehlo.add %15960, %15963 : tensor<1x64xf32> loc(#loc5263)
    %15965 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15966 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15967 = stablehlo.broadcast_in_dim %15966, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15968 = stablehlo.add %15965, %15967 : tensor<1xf32> loc(#loc5259)
    %15969 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15970 = stablehlo.broadcast_in_dim %15969, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15971 = stablehlo.subtract %15968, %15970 : tensor<1xf32> loc(#loc5260)
    %15972 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15973 = stablehlo.broadcast_in_dim %15972, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15974 = stablehlo.power %15973, %15971 : tensor<1xf32> loc(#loc5261)
    %15975 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15976 = stablehlo.broadcast_in_dim %15975, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15977 = stablehlo.subtract %15976, %15974 : tensor<1xf32> loc(#loc5260)
    %15978 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15979 = stablehlo.broadcast_in_dim %15978, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %15980 = stablehlo.multiply %15979, %15977 : tensor<1xf32> loc(#loc5236)
    %15981 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %15982 = stablehlo.broadcast_in_dim %15981, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15983 = stablehlo.power %15982, %15968 : tensor<1xf32> loc(#loc5261)
    %15984 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15985 = stablehlo.broadcast_in_dim %15984, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15986 = stablehlo.subtract %15985, %15983 : tensor<1xf32> loc(#loc5260)
    %15987 = stablehlo.divide %15980, %15986 : tensor<1xf32> loc(#loc5253)
    %15988 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %15989 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15990 = stablehlo.broadcast_in_dim %15989, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %15991 = stablehlo.add %15988, %15990 : tensor<1xf32> loc(#loc5259)
    %15992 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15993 = stablehlo.broadcast_in_dim %15992, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %15994 = stablehlo.subtract %15991, %15993 : tensor<1xf32> loc(#loc5260)
    %15995 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %15996 = stablehlo.broadcast_in_dim %15995, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %15997 = stablehlo.power %15996, %15994 : tensor<1xf32> loc(#loc5261)
    %15998 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %15999 = stablehlo.broadcast_in_dim %15998, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16000 = stablehlo.subtract %15999, %15997 : tensor<1xf32> loc(#loc5260)
    %16001 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16002 = stablehlo.broadcast_in_dim %16001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16003 = stablehlo.multiply %16002, %16000 : tensor<1xf32> loc(#loc5236)
    %16004 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16005 = stablehlo.broadcast_in_dim %16004, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16006 = stablehlo.power %16005, %15991 : tensor<1xf32> loc(#loc5261)
    %16007 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16008 = stablehlo.broadcast_in_dim %16007, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16009 = stablehlo.subtract %16008, %16006 : tensor<1xf32> loc(#loc5260)
    %16010 = stablehlo.divide %16003, %16009 : tensor<1xf32> loc(#loc5253)
    %16011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16012 = stablehlo.broadcast_in_dim %16011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16013 = stablehlo.subtract %16012, %15987 : tensor<1xf32> loc(#loc5260)
    %16014 = stablehlo.broadcast_in_dim %16013, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16015 = stablehlo.broadcast_in_dim %16014, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16016 = stablehlo.multiply %16015, %4932 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16017 = stablehlo.broadcast_in_dim %15987, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16018 = stablehlo.broadcast_in_dim %16017, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16019 = stablehlo.multiply %16018, %arg363 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16020 = stablehlo.add %16016, %16019 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %16021 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16022 = stablehlo.broadcast_in_dim %16021, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16023 = stablehlo.subtract %16022, %16010 : tensor<1xf32> loc(#loc5260)
    %16024 = stablehlo.multiply %4932, %4932 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %16025 = stablehlo.broadcast_in_dim %16023, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16026 = stablehlo.broadcast_in_dim %16025, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16027 = stablehlo.multiply %16026, %16024 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16028 = stablehlo.broadcast_in_dim %16010, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16029 = stablehlo.broadcast_in_dim %16028, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16030 = stablehlo.multiply %16029, %arg539 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16031 = stablehlo.add %16027, %16030 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %16032 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16033 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16034 = stablehlo.broadcast_in_dim %16033, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16035 = stablehlo.add %16032, %16034 : tensor<1xf32> loc(#loc5259)
    %16036 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16037 = stablehlo.broadcast_in_dim %16036, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16038 = stablehlo.subtract %16035, %16037 : tensor<1xf32> loc(#loc5260)
    %16039 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16040 = stablehlo.broadcast_in_dim %16039, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16041 = stablehlo.power %16040, %16038 : tensor<1xf32> loc(#loc5261)
    %16042 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16043 = stablehlo.broadcast_in_dim %16042, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16044 = stablehlo.subtract %16043, %16041 : tensor<1xf32> loc(#loc5260)
    %16045 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16046 = stablehlo.broadcast_in_dim %16045, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16047 = stablehlo.multiply %16046, %16044 : tensor<1xf32> loc(#loc5236)
    %16048 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16049 = stablehlo.broadcast_in_dim %16048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16050 = stablehlo.power %16049, %16035 : tensor<1xf32> loc(#loc5261)
    %16051 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16052 = stablehlo.broadcast_in_dim %16051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16053 = stablehlo.subtract %16052, %16050 : tensor<1xf32> loc(#loc5260)
    %16054 = stablehlo.divide %16047, %16053 : tensor<1xf32> loc(#loc5253)
    %16055 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16056 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16057 = stablehlo.broadcast_in_dim %16056, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16058 = stablehlo.add %16055, %16057 : tensor<1xf32> loc(#loc5259)
    %16059 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16060 = stablehlo.broadcast_in_dim %16059, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16061 = stablehlo.subtract %16058, %16060 : tensor<1xf32> loc(#loc5260)
    %16062 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16063 = stablehlo.broadcast_in_dim %16062, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16064 = stablehlo.power %16063, %16061 : tensor<1xf32> loc(#loc5261)
    %16065 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16066 = stablehlo.broadcast_in_dim %16065, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16067 = stablehlo.subtract %16066, %16064 : tensor<1xf32> loc(#loc5260)
    %16068 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16069 = stablehlo.broadcast_in_dim %16068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16070 = stablehlo.multiply %16069, %16067 : tensor<1xf32> loc(#loc5236)
    %16071 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16072 = stablehlo.broadcast_in_dim %16071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16073 = stablehlo.power %16072, %16058 : tensor<1xf32> loc(#loc5261)
    %16074 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16075 = stablehlo.broadcast_in_dim %16074, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16076 = stablehlo.subtract %16075, %16073 : tensor<1xf32> loc(#loc5260)
    %16077 = stablehlo.divide %16070, %16076 : tensor<1xf32> loc(#loc5253)
    %16078 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16079 = stablehlo.broadcast_in_dim %16078, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16080 = stablehlo.subtract %16079, %16054 : tensor<1xf32> loc(#loc5260)
    %16081 = stablehlo.broadcast_in_dim %16080, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16082 = stablehlo.broadcast_in_dim %16081, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %16083 = stablehlo.multiply %16082, %4936 : tensor<1x8192xf32> loc(#loc5262)
    %16084 = stablehlo.broadcast_in_dim %16054, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16085 = stablehlo.broadcast_in_dim %16084, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %16086 = stablehlo.multiply %16085, %arg364 : tensor<1x8192xf32> loc(#loc5262)
    %16087 = stablehlo.add %16083, %16086 : tensor<1x8192xf32> loc(#loc5263)
    %16088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16089 = stablehlo.broadcast_in_dim %16088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16090 = stablehlo.subtract %16089, %16077 : tensor<1xf32> loc(#loc5260)
    %16091 = stablehlo.multiply %4936, %4936 : tensor<1x8192xf32> loc(#loc5264)
    %16092 = stablehlo.broadcast_in_dim %16090, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16093 = stablehlo.broadcast_in_dim %16092, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %16094 = stablehlo.multiply %16093, %16091 : tensor<1x8192xf32> loc(#loc5262)
    %16095 = stablehlo.broadcast_in_dim %16077, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16096 = stablehlo.broadcast_in_dim %16095, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %16097 = stablehlo.multiply %16096, %arg540 : tensor<1x8192xf32> loc(#loc5262)
    %16098 = stablehlo.add %16094, %16097 : tensor<1x8192xf32> loc(#loc5263)
    %16099 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16100 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16101 = stablehlo.broadcast_in_dim %16100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16102 = stablehlo.add %16099, %16101 : tensor<1xf32> loc(#loc5259)
    %16103 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16104 = stablehlo.broadcast_in_dim %16103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16105 = stablehlo.subtract %16102, %16104 : tensor<1xf32> loc(#loc5260)
    %16106 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16107 = stablehlo.broadcast_in_dim %16106, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16108 = stablehlo.power %16107, %16105 : tensor<1xf32> loc(#loc5261)
    %16109 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16110 = stablehlo.broadcast_in_dim %16109, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16111 = stablehlo.subtract %16110, %16108 : tensor<1xf32> loc(#loc5260)
    %16112 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16113 = stablehlo.broadcast_in_dim %16112, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16114 = stablehlo.multiply %16113, %16111 : tensor<1xf32> loc(#loc5236)
    %16115 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16116 = stablehlo.broadcast_in_dim %16115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16117 = stablehlo.power %16116, %16102 : tensor<1xf32> loc(#loc5261)
    %16118 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16119 = stablehlo.broadcast_in_dim %16118, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16120 = stablehlo.subtract %16119, %16117 : tensor<1xf32> loc(#loc5260)
    %16121 = stablehlo.divide %16114, %16120 : tensor<1xf32> loc(#loc5253)
    %16122 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16123 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16124 = stablehlo.broadcast_in_dim %16123, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16125 = stablehlo.add %16122, %16124 : tensor<1xf32> loc(#loc5259)
    %16126 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16127 = stablehlo.broadcast_in_dim %16126, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16128 = stablehlo.subtract %16125, %16127 : tensor<1xf32> loc(#loc5260)
    %16129 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16130 = stablehlo.broadcast_in_dim %16129, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16131 = stablehlo.power %16130, %16128 : tensor<1xf32> loc(#loc5261)
    %16132 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16133 = stablehlo.broadcast_in_dim %16132, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16134 = stablehlo.subtract %16133, %16131 : tensor<1xf32> loc(#loc5260)
    %16135 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16136 = stablehlo.broadcast_in_dim %16135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16137 = stablehlo.multiply %16136, %16134 : tensor<1xf32> loc(#loc5236)
    %16138 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16139 = stablehlo.broadcast_in_dim %16138, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16140 = stablehlo.power %16139, %16125 : tensor<1xf32> loc(#loc5261)
    %16141 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16142 = stablehlo.broadcast_in_dim %16141, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16143 = stablehlo.subtract %16142, %16140 : tensor<1xf32> loc(#loc5260)
    %16144 = stablehlo.divide %16137, %16143 : tensor<1xf32> loc(#loc5253)
    %16145 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16146 = stablehlo.broadcast_in_dim %16145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16147 = stablehlo.subtract %16146, %16121 : tensor<1xf32> loc(#loc5260)
    %16148 = stablehlo.broadcast_in_dim %16147, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16149 = stablehlo.broadcast_in_dim %16148, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %16150 = stablehlo.multiply %16149, %4940 : tensor<1x2048x8192xf32> loc(#loc5262)
    %16151 = stablehlo.broadcast_in_dim %16121, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16152 = stablehlo.broadcast_in_dim %16151, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %16153 = stablehlo.multiply %16152, %arg365 : tensor<1x2048x8192xf32> loc(#loc5262)
    %16154 = stablehlo.add %16150, %16153 : tensor<1x2048x8192xf32> loc(#loc5263)
    %16155 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16156 = stablehlo.broadcast_in_dim %16155, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16157 = stablehlo.subtract %16156, %16144 : tensor<1xf32> loc(#loc5260)
    %16158 = stablehlo.multiply %4940, %4940 : tensor<1x2048x8192xf32> loc(#loc5264)
    %16159 = stablehlo.broadcast_in_dim %16157, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16160 = stablehlo.broadcast_in_dim %16159, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %16161 = stablehlo.multiply %16160, %16158 : tensor<1x2048x8192xf32> loc(#loc5262)
    %16162 = stablehlo.broadcast_in_dim %16144, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16163 = stablehlo.broadcast_in_dim %16162, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %16164 = stablehlo.multiply %16163, %arg541 : tensor<1x2048x8192xf32> loc(#loc5262)
    %16165 = stablehlo.add %16161, %16164 : tensor<1x2048x8192xf32> loc(#loc5263)
    %16166 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16167 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16168 = stablehlo.broadcast_in_dim %16167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16169 = stablehlo.add %16166, %16168 : tensor<1xf32> loc(#loc5259)
    %16170 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16171 = stablehlo.broadcast_in_dim %16170, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16172 = stablehlo.subtract %16169, %16171 : tensor<1xf32> loc(#loc5260)
    %16173 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16174 = stablehlo.broadcast_in_dim %16173, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16175 = stablehlo.power %16174, %16172 : tensor<1xf32> loc(#loc5261)
    %16176 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16177 = stablehlo.broadcast_in_dim %16176, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16178 = stablehlo.subtract %16177, %16175 : tensor<1xf32> loc(#loc5260)
    %16179 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16180 = stablehlo.broadcast_in_dim %16179, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16181 = stablehlo.multiply %16180, %16178 : tensor<1xf32> loc(#loc5236)
    %16182 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16183 = stablehlo.broadcast_in_dim %16182, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16184 = stablehlo.power %16183, %16169 : tensor<1xf32> loc(#loc5261)
    %16185 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16186 = stablehlo.broadcast_in_dim %16185, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16187 = stablehlo.subtract %16186, %16184 : tensor<1xf32> loc(#loc5260)
    %16188 = stablehlo.divide %16181, %16187 : tensor<1xf32> loc(#loc5253)
    %16189 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16190 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16191 = stablehlo.broadcast_in_dim %16190, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16192 = stablehlo.add %16189, %16191 : tensor<1xf32> loc(#loc5259)
    %16193 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16194 = stablehlo.broadcast_in_dim %16193, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16195 = stablehlo.subtract %16192, %16194 : tensor<1xf32> loc(#loc5260)
    %16196 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16197 = stablehlo.broadcast_in_dim %16196, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16198 = stablehlo.power %16197, %16195 : tensor<1xf32> loc(#loc5261)
    %16199 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16200 = stablehlo.broadcast_in_dim %16199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16201 = stablehlo.subtract %16200, %16198 : tensor<1xf32> loc(#loc5260)
    %16202 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16203 = stablehlo.broadcast_in_dim %16202, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16204 = stablehlo.multiply %16203, %16201 : tensor<1xf32> loc(#loc5236)
    %16205 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16206 = stablehlo.broadcast_in_dim %16205, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16207 = stablehlo.power %16206, %16192 : tensor<1xf32> loc(#loc5261)
    %16208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16209 = stablehlo.broadcast_in_dim %16208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16210 = stablehlo.subtract %16209, %16207 : tensor<1xf32> loc(#loc5260)
    %16211 = stablehlo.divide %16204, %16210 : tensor<1xf32> loc(#loc5253)
    %16212 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16213 = stablehlo.broadcast_in_dim %16212, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16214 = stablehlo.subtract %16213, %16188 : tensor<1xf32> loc(#loc5260)
    %16215 = stablehlo.broadcast_in_dim %16214, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16216 = stablehlo.broadcast_in_dim %16215, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16217 = stablehlo.multiply %16216, %4944 : tensor<1x2048xf32> loc(#loc5262)
    %16218 = stablehlo.broadcast_in_dim %16188, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16219 = stablehlo.broadcast_in_dim %16218, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16220 = stablehlo.multiply %16219, %arg366 : tensor<1x2048xf32> loc(#loc5262)
    %16221 = stablehlo.add %16217, %16220 : tensor<1x2048xf32> loc(#loc5263)
    %16222 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16223 = stablehlo.broadcast_in_dim %16222, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16224 = stablehlo.subtract %16223, %16211 : tensor<1xf32> loc(#loc5260)
    %16225 = stablehlo.multiply %4944, %4944 : tensor<1x2048xf32> loc(#loc5264)
    %16226 = stablehlo.broadcast_in_dim %16224, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16227 = stablehlo.broadcast_in_dim %16226, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16228 = stablehlo.multiply %16227, %16225 : tensor<1x2048xf32> loc(#loc5262)
    %16229 = stablehlo.broadcast_in_dim %16211, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16230 = stablehlo.broadcast_in_dim %16229, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16231 = stablehlo.multiply %16230, %arg542 : tensor<1x2048xf32> loc(#loc5262)
    %16232 = stablehlo.add %16228, %16231 : tensor<1x2048xf32> loc(#loc5263)
    %16233 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16234 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16235 = stablehlo.broadcast_in_dim %16234, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16236 = stablehlo.add %16233, %16235 : tensor<1xf32> loc(#loc5259)
    %16237 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16238 = stablehlo.broadcast_in_dim %16237, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16239 = stablehlo.subtract %16236, %16238 : tensor<1xf32> loc(#loc5260)
    %16240 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16241 = stablehlo.broadcast_in_dim %16240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16242 = stablehlo.power %16241, %16239 : tensor<1xf32> loc(#loc5261)
    %16243 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16244 = stablehlo.broadcast_in_dim %16243, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16245 = stablehlo.subtract %16244, %16242 : tensor<1xf32> loc(#loc5260)
    %16246 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16247 = stablehlo.broadcast_in_dim %16246, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16248 = stablehlo.multiply %16247, %16245 : tensor<1xf32> loc(#loc5236)
    %16249 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16250 = stablehlo.broadcast_in_dim %16249, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16251 = stablehlo.power %16250, %16236 : tensor<1xf32> loc(#loc5261)
    %16252 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16253 = stablehlo.broadcast_in_dim %16252, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16254 = stablehlo.subtract %16253, %16251 : tensor<1xf32> loc(#loc5260)
    %16255 = stablehlo.divide %16248, %16254 : tensor<1xf32> loc(#loc5253)
    %16256 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16257 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16258 = stablehlo.broadcast_in_dim %16257, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16259 = stablehlo.add %16256, %16258 : tensor<1xf32> loc(#loc5259)
    %16260 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16261 = stablehlo.broadcast_in_dim %16260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16262 = stablehlo.subtract %16259, %16261 : tensor<1xf32> loc(#loc5260)
    %16263 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16264 = stablehlo.broadcast_in_dim %16263, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16265 = stablehlo.power %16264, %16262 : tensor<1xf32> loc(#loc5261)
    %16266 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16267 = stablehlo.broadcast_in_dim %16266, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16268 = stablehlo.subtract %16267, %16265 : tensor<1xf32> loc(#loc5260)
    %16269 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16270 = stablehlo.broadcast_in_dim %16269, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16271 = stablehlo.multiply %16270, %16268 : tensor<1xf32> loc(#loc5236)
    %16272 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16273 = stablehlo.broadcast_in_dim %16272, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16274 = stablehlo.power %16273, %16259 : tensor<1xf32> loc(#loc5261)
    %16275 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16276 = stablehlo.broadcast_in_dim %16275, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16277 = stablehlo.subtract %16276, %16274 : tensor<1xf32> loc(#loc5260)
    %16278 = stablehlo.divide %16271, %16277 : tensor<1xf32> loc(#loc5253)
    %16279 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16280 = stablehlo.broadcast_in_dim %16279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16281 = stablehlo.subtract %16280, %16255 : tensor<1xf32> loc(#loc5260)
    %16282 = stablehlo.broadcast_in_dim %16281, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16283 = stablehlo.broadcast_in_dim %16282, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %16284 = stablehlo.multiply %16283, %4948 : tensor<1x8192x2048xf32> loc(#loc5262)
    %16285 = stablehlo.broadcast_in_dim %16255, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16286 = stablehlo.broadcast_in_dim %16285, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %16287 = stablehlo.multiply %16286, %arg367 : tensor<1x8192x2048xf32> loc(#loc5262)
    %16288 = stablehlo.add %16284, %16287 : tensor<1x8192x2048xf32> loc(#loc5263)
    %16289 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16290 = stablehlo.broadcast_in_dim %16289, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16291 = stablehlo.subtract %16290, %16278 : tensor<1xf32> loc(#loc5260)
    %16292 = stablehlo.multiply %4948, %4948 : tensor<1x8192x2048xf32> loc(#loc5264)
    %16293 = stablehlo.broadcast_in_dim %16291, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16294 = stablehlo.broadcast_in_dim %16293, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %16295 = stablehlo.multiply %16294, %16292 : tensor<1x8192x2048xf32> loc(#loc5262)
    %16296 = stablehlo.broadcast_in_dim %16278, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %16297 = stablehlo.broadcast_in_dim %16296, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %16298 = stablehlo.multiply %16297, %arg543 : tensor<1x8192x2048xf32> loc(#loc5262)
    %16299 = stablehlo.add %16295, %16298 : tensor<1x8192x2048xf32> loc(#loc5263)
    %16300 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16301 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16302 = stablehlo.broadcast_in_dim %16301, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16303 = stablehlo.add %16300, %16302 : tensor<1xf32> loc(#loc5259)
    %16304 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16305 = stablehlo.broadcast_in_dim %16304, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16306 = stablehlo.subtract %16303, %16305 : tensor<1xf32> loc(#loc5260)
    %16307 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16308 = stablehlo.broadcast_in_dim %16307, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16309 = stablehlo.power %16308, %16306 : tensor<1xf32> loc(#loc5261)
    %16310 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16311 = stablehlo.broadcast_in_dim %16310, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16312 = stablehlo.subtract %16311, %16309 : tensor<1xf32> loc(#loc5260)
    %16313 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16314 = stablehlo.broadcast_in_dim %16313, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16315 = stablehlo.multiply %16314, %16312 : tensor<1xf32> loc(#loc5236)
    %16316 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16317 = stablehlo.broadcast_in_dim %16316, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16318 = stablehlo.power %16317, %16303 : tensor<1xf32> loc(#loc5261)
    %16319 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16320 = stablehlo.broadcast_in_dim %16319, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16321 = stablehlo.subtract %16320, %16318 : tensor<1xf32> loc(#loc5260)
    %16322 = stablehlo.divide %16315, %16321 : tensor<1xf32> loc(#loc5253)
    %16323 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16324 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16325 = stablehlo.broadcast_in_dim %16324, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16326 = stablehlo.add %16323, %16325 : tensor<1xf32> loc(#loc5259)
    %16327 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16328 = stablehlo.broadcast_in_dim %16327, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16329 = stablehlo.subtract %16326, %16328 : tensor<1xf32> loc(#loc5260)
    %16330 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16331 = stablehlo.broadcast_in_dim %16330, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16332 = stablehlo.power %16331, %16329 : tensor<1xf32> loc(#loc5261)
    %16333 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16334 = stablehlo.broadcast_in_dim %16333, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16335 = stablehlo.subtract %16334, %16332 : tensor<1xf32> loc(#loc5260)
    %16336 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16337 = stablehlo.broadcast_in_dim %16336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16338 = stablehlo.multiply %16337, %16335 : tensor<1xf32> loc(#loc5236)
    %16339 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16340 = stablehlo.broadcast_in_dim %16339, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16341 = stablehlo.power %16340, %16326 : tensor<1xf32> loc(#loc5261)
    %16342 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16343 = stablehlo.broadcast_in_dim %16342, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16344 = stablehlo.subtract %16343, %16341 : tensor<1xf32> loc(#loc5260)
    %16345 = stablehlo.divide %16338, %16344 : tensor<1xf32> loc(#loc5253)
    %16346 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16347 = stablehlo.broadcast_in_dim %16346, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16348 = stablehlo.subtract %16347, %16322 : tensor<1xf32> loc(#loc5260)
    %16349 = stablehlo.broadcast_in_dim %16348, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16350 = stablehlo.broadcast_in_dim %16349, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16351 = stablehlo.multiply %16350, %4952 : tensor<1x2048xf32> loc(#loc5262)
    %16352 = stablehlo.broadcast_in_dim %16322, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16353 = stablehlo.broadcast_in_dim %16352, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16354 = stablehlo.multiply %16353, %arg368 : tensor<1x2048xf32> loc(#loc5262)
    %16355 = stablehlo.add %16351, %16354 : tensor<1x2048xf32> loc(#loc5263)
    %16356 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16357 = stablehlo.broadcast_in_dim %16356, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16358 = stablehlo.subtract %16357, %16345 : tensor<1xf32> loc(#loc5260)
    %16359 = stablehlo.multiply %4952, %4952 : tensor<1x2048xf32> loc(#loc5264)
    %16360 = stablehlo.broadcast_in_dim %16358, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16361 = stablehlo.broadcast_in_dim %16360, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16362 = stablehlo.multiply %16361, %16359 : tensor<1x2048xf32> loc(#loc5262)
    %16363 = stablehlo.broadcast_in_dim %16345, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16364 = stablehlo.broadcast_in_dim %16363, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16365 = stablehlo.multiply %16364, %arg544 : tensor<1x2048xf32> loc(#loc5262)
    %16366 = stablehlo.add %16362, %16365 : tensor<1x2048xf32> loc(#loc5263)
    %16367 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16369 = stablehlo.broadcast_in_dim %16368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16370 = stablehlo.add %16367, %16369 : tensor<1xf32> loc(#loc5259)
    %16371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16372 = stablehlo.broadcast_in_dim %16371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16373 = stablehlo.subtract %16370, %16372 : tensor<1xf32> loc(#loc5260)
    %16374 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16375 = stablehlo.broadcast_in_dim %16374, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16376 = stablehlo.power %16375, %16373 : tensor<1xf32> loc(#loc5261)
    %16377 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16378 = stablehlo.broadcast_in_dim %16377, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16379 = stablehlo.subtract %16378, %16376 : tensor<1xf32> loc(#loc5260)
    %16380 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16381 = stablehlo.broadcast_in_dim %16380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16382 = stablehlo.multiply %16381, %16379 : tensor<1xf32> loc(#loc5236)
    %16383 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16384 = stablehlo.broadcast_in_dim %16383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16385 = stablehlo.power %16384, %16370 : tensor<1xf32> loc(#loc5261)
    %16386 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16387 = stablehlo.broadcast_in_dim %16386, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16388 = stablehlo.subtract %16387, %16385 : tensor<1xf32> loc(#loc5260)
    %16389 = stablehlo.divide %16382, %16388 : tensor<1xf32> loc(#loc5253)
    %16390 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16391 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16392 = stablehlo.broadcast_in_dim %16391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16393 = stablehlo.add %16390, %16392 : tensor<1xf32> loc(#loc5259)
    %16394 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16395 = stablehlo.broadcast_in_dim %16394, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16396 = stablehlo.subtract %16393, %16395 : tensor<1xf32> loc(#loc5260)
    %16397 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16398 = stablehlo.broadcast_in_dim %16397, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16399 = stablehlo.power %16398, %16396 : tensor<1xf32> loc(#loc5261)
    %16400 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16401 = stablehlo.broadcast_in_dim %16400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16402 = stablehlo.subtract %16401, %16399 : tensor<1xf32> loc(#loc5260)
    %16403 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16404 = stablehlo.broadcast_in_dim %16403, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16405 = stablehlo.multiply %16404, %16402 : tensor<1xf32> loc(#loc5236)
    %16406 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16407 = stablehlo.broadcast_in_dim %16406, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16408 = stablehlo.power %16407, %16393 : tensor<1xf32> loc(#loc5261)
    %16409 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16410 = stablehlo.broadcast_in_dim %16409, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16411 = stablehlo.subtract %16410, %16408 : tensor<1xf32> loc(#loc5260)
    %16412 = stablehlo.divide %16405, %16411 : tensor<1xf32> loc(#loc5253)
    %16413 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16414 = stablehlo.broadcast_in_dim %16413, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16415 = stablehlo.subtract %16414, %16389 : tensor<1xf32> loc(#loc5260)
    %16416 = stablehlo.broadcast_in_dim %16415, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16417 = stablehlo.broadcast_in_dim %16416, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16418 = stablehlo.multiply %16417, %4956 : tensor<1x2048xf32> loc(#loc5262)
    %16419 = stablehlo.broadcast_in_dim %16389, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16420 = stablehlo.broadcast_in_dim %16419, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16421 = stablehlo.multiply %16420, %arg369 : tensor<1x2048xf32> loc(#loc5262)
    %16422 = stablehlo.add %16418, %16421 : tensor<1x2048xf32> loc(#loc5263)
    %16423 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16424 = stablehlo.broadcast_in_dim %16423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16425 = stablehlo.subtract %16424, %16412 : tensor<1xf32> loc(#loc5260)
    %16426 = stablehlo.multiply %4956, %4956 : tensor<1x2048xf32> loc(#loc5264)
    %16427 = stablehlo.broadcast_in_dim %16425, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16428 = stablehlo.broadcast_in_dim %16427, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16429 = stablehlo.multiply %16428, %16426 : tensor<1x2048xf32> loc(#loc5262)
    %16430 = stablehlo.broadcast_in_dim %16412, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16431 = stablehlo.broadcast_in_dim %16430, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16432 = stablehlo.multiply %16431, %arg545 : tensor<1x2048xf32> loc(#loc5262)
    %16433 = stablehlo.add %16429, %16432 : tensor<1x2048xf32> loc(#loc5263)
    %16434 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16435 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16436 = stablehlo.broadcast_in_dim %16435, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16437 = stablehlo.add %16434, %16436 : tensor<1xf32> loc(#loc5259)
    %16438 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16439 = stablehlo.broadcast_in_dim %16438, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16440 = stablehlo.subtract %16437, %16439 : tensor<1xf32> loc(#loc5260)
    %16441 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16442 = stablehlo.broadcast_in_dim %16441, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16443 = stablehlo.power %16442, %16440 : tensor<1xf32> loc(#loc5261)
    %16444 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16445 = stablehlo.broadcast_in_dim %16444, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16446 = stablehlo.subtract %16445, %16443 : tensor<1xf32> loc(#loc5260)
    %16447 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16448 = stablehlo.broadcast_in_dim %16447, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16449 = stablehlo.multiply %16448, %16446 : tensor<1xf32> loc(#loc5236)
    %16450 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16451 = stablehlo.broadcast_in_dim %16450, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16452 = stablehlo.power %16451, %16437 : tensor<1xf32> loc(#loc5261)
    %16453 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16454 = stablehlo.broadcast_in_dim %16453, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16455 = stablehlo.subtract %16454, %16452 : tensor<1xf32> loc(#loc5260)
    %16456 = stablehlo.divide %16449, %16455 : tensor<1xf32> loc(#loc5253)
    %16457 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16458 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16459 = stablehlo.broadcast_in_dim %16458, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16460 = stablehlo.add %16457, %16459 : tensor<1xf32> loc(#loc5259)
    %16461 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16462 = stablehlo.broadcast_in_dim %16461, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16463 = stablehlo.subtract %16460, %16462 : tensor<1xf32> loc(#loc5260)
    %16464 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16465 = stablehlo.broadcast_in_dim %16464, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16466 = stablehlo.power %16465, %16463 : tensor<1xf32> loc(#loc5261)
    %16467 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16468 = stablehlo.broadcast_in_dim %16467, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16469 = stablehlo.subtract %16468, %16466 : tensor<1xf32> loc(#loc5260)
    %16470 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16471 = stablehlo.broadcast_in_dim %16470, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16472 = stablehlo.multiply %16471, %16469 : tensor<1xf32> loc(#loc5236)
    %16473 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16474 = stablehlo.broadcast_in_dim %16473, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16475 = stablehlo.power %16474, %16460 : tensor<1xf32> loc(#loc5261)
    %16476 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16477 = stablehlo.broadcast_in_dim %16476, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16478 = stablehlo.subtract %16477, %16475 : tensor<1xf32> loc(#loc5260)
    %16479 = stablehlo.divide %16472, %16478 : tensor<1xf32> loc(#loc5253)
    %16480 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16481 = stablehlo.broadcast_in_dim %16480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16482 = stablehlo.subtract %16481, %16456 : tensor<1xf32> loc(#loc5260)
    %16483 = stablehlo.broadcast_in_dim %16482, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16484 = stablehlo.broadcast_in_dim %16483, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16485 = stablehlo.multiply %16484, %4960 : tensor<1x2048xf32> loc(#loc5262)
    %16486 = stablehlo.broadcast_in_dim %16456, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16487 = stablehlo.broadcast_in_dim %16486, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16488 = stablehlo.multiply %16487, %arg370 : tensor<1x2048xf32> loc(#loc5262)
    %16489 = stablehlo.add %16485, %16488 : tensor<1x2048xf32> loc(#loc5263)
    %16490 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16491 = stablehlo.broadcast_in_dim %16490, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16492 = stablehlo.subtract %16491, %16479 : tensor<1xf32> loc(#loc5260)
    %16493 = stablehlo.multiply %4960, %4960 : tensor<1x2048xf32> loc(#loc5264)
    %16494 = stablehlo.broadcast_in_dim %16492, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16495 = stablehlo.broadcast_in_dim %16494, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16496 = stablehlo.multiply %16495, %16493 : tensor<1x2048xf32> loc(#loc5262)
    %16497 = stablehlo.broadcast_in_dim %16479, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16498 = stablehlo.broadcast_in_dim %16497, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16499 = stablehlo.multiply %16498, %arg546 : tensor<1x2048xf32> loc(#loc5262)
    %16500 = stablehlo.add %16496, %16499 : tensor<1x2048xf32> loc(#loc5263)
    %16501 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16502 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16503 = stablehlo.broadcast_in_dim %16502, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16504 = stablehlo.add %16501, %16503 : tensor<1xf32> loc(#loc5259)
    %16505 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16506 = stablehlo.broadcast_in_dim %16505, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16507 = stablehlo.subtract %16504, %16506 : tensor<1xf32> loc(#loc5260)
    %16508 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16509 = stablehlo.broadcast_in_dim %16508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16510 = stablehlo.power %16509, %16507 : tensor<1xf32> loc(#loc5261)
    %16511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16512 = stablehlo.broadcast_in_dim %16511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16513 = stablehlo.subtract %16512, %16510 : tensor<1xf32> loc(#loc5260)
    %16514 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16515 = stablehlo.broadcast_in_dim %16514, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16516 = stablehlo.multiply %16515, %16513 : tensor<1xf32> loc(#loc5236)
    %16517 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16518 = stablehlo.broadcast_in_dim %16517, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16519 = stablehlo.power %16518, %16504 : tensor<1xf32> loc(#loc5261)
    %16520 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16521 = stablehlo.broadcast_in_dim %16520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16522 = stablehlo.subtract %16521, %16519 : tensor<1xf32> loc(#loc5260)
    %16523 = stablehlo.divide %16516, %16522 : tensor<1xf32> loc(#loc5253)
    %16524 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16525 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16526 = stablehlo.broadcast_in_dim %16525, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16527 = stablehlo.add %16524, %16526 : tensor<1xf32> loc(#loc5259)
    %16528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16529 = stablehlo.broadcast_in_dim %16528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16530 = stablehlo.subtract %16527, %16529 : tensor<1xf32> loc(#loc5260)
    %16531 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16532 = stablehlo.broadcast_in_dim %16531, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16533 = stablehlo.power %16532, %16530 : tensor<1xf32> loc(#loc5261)
    %16534 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16535 = stablehlo.broadcast_in_dim %16534, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16536 = stablehlo.subtract %16535, %16533 : tensor<1xf32> loc(#loc5260)
    %16537 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16538 = stablehlo.broadcast_in_dim %16537, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16539 = stablehlo.multiply %16538, %16536 : tensor<1xf32> loc(#loc5236)
    %16540 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16541 = stablehlo.broadcast_in_dim %16540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16542 = stablehlo.power %16541, %16527 : tensor<1xf32> loc(#loc5261)
    %16543 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16544 = stablehlo.broadcast_in_dim %16543, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16545 = stablehlo.subtract %16544, %16542 : tensor<1xf32> loc(#loc5260)
    %16546 = stablehlo.divide %16539, %16545 : tensor<1xf32> loc(#loc5253)
    %16547 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16548 = stablehlo.broadcast_in_dim %16547, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16549 = stablehlo.subtract %16548, %16523 : tensor<1xf32> loc(#loc5260)
    %16550 = stablehlo.broadcast_in_dim %16549, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16551 = stablehlo.broadcast_in_dim %16550, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16552 = stablehlo.multiply %16551, %4964 : tensor<1x2048xf32> loc(#loc5262)
    %16553 = stablehlo.broadcast_in_dim %16523, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16554 = stablehlo.broadcast_in_dim %16553, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16555 = stablehlo.multiply %16554, %arg371 : tensor<1x2048xf32> loc(#loc5262)
    %16556 = stablehlo.add %16552, %16555 : tensor<1x2048xf32> loc(#loc5263)
    %16557 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16558 = stablehlo.broadcast_in_dim %16557, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16559 = stablehlo.subtract %16558, %16546 : tensor<1xf32> loc(#loc5260)
    %16560 = stablehlo.multiply %4964, %4964 : tensor<1x2048xf32> loc(#loc5264)
    %16561 = stablehlo.broadcast_in_dim %16559, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16562 = stablehlo.broadcast_in_dim %16561, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16563 = stablehlo.multiply %16562, %16560 : tensor<1x2048xf32> loc(#loc5262)
    %16564 = stablehlo.broadcast_in_dim %16546, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16565 = stablehlo.broadcast_in_dim %16564, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %16566 = stablehlo.multiply %16565, %arg547 : tensor<1x2048xf32> loc(#loc5262)
    %16567 = stablehlo.add %16563, %16566 : tensor<1x2048xf32> loc(#loc5263)
    %16568 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16569 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16570 = stablehlo.broadcast_in_dim %16569, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16571 = stablehlo.add %16568, %16570 : tensor<1xf32> loc(#loc5259)
    %16572 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16573 = stablehlo.broadcast_in_dim %16572, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16574 = stablehlo.subtract %16571, %16573 : tensor<1xf32> loc(#loc5260)
    %16575 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16576 = stablehlo.broadcast_in_dim %16575, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16577 = stablehlo.power %16576, %16574 : tensor<1xf32> loc(#loc5261)
    %16578 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16579 = stablehlo.broadcast_in_dim %16578, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16580 = stablehlo.subtract %16579, %16577 : tensor<1xf32> loc(#loc5260)
    %16581 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16582 = stablehlo.broadcast_in_dim %16581, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16583 = stablehlo.multiply %16582, %16580 : tensor<1xf32> loc(#loc5236)
    %16584 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16585 = stablehlo.broadcast_in_dim %16584, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16586 = stablehlo.power %16585, %16571 : tensor<1xf32> loc(#loc5261)
    %16587 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16588 = stablehlo.broadcast_in_dim %16587, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16589 = stablehlo.subtract %16588, %16586 : tensor<1xf32> loc(#loc5260)
    %16590 = stablehlo.divide %16583, %16589 : tensor<1xf32> loc(#loc5253)
    %16591 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16592 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16593 = stablehlo.broadcast_in_dim %16592, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16594 = stablehlo.add %16591, %16593 : tensor<1xf32> loc(#loc5259)
    %16595 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16596 = stablehlo.broadcast_in_dim %16595, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16597 = stablehlo.subtract %16594, %16596 : tensor<1xf32> loc(#loc5260)
    %16598 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16599 = stablehlo.broadcast_in_dim %16598, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16600 = stablehlo.power %16599, %16597 : tensor<1xf32> loc(#loc5261)
    %16601 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16602 = stablehlo.broadcast_in_dim %16601, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16603 = stablehlo.subtract %16602, %16600 : tensor<1xf32> loc(#loc5260)
    %16604 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16605 = stablehlo.broadcast_in_dim %16604, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16606 = stablehlo.multiply %16605, %16603 : tensor<1xf32> loc(#loc5236)
    %16607 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16608 = stablehlo.broadcast_in_dim %16607, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16609 = stablehlo.power %16608, %16594 : tensor<1xf32> loc(#loc5261)
    %16610 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16611 = stablehlo.broadcast_in_dim %16610, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16612 = stablehlo.subtract %16611, %16609 : tensor<1xf32> loc(#loc5260)
    %16613 = stablehlo.divide %16606, %16612 : tensor<1xf32> loc(#loc5253)
    %16614 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16615 = stablehlo.broadcast_in_dim %16614, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16616 = stablehlo.subtract %16615, %16590 : tensor<1xf32> loc(#loc5260)
    %16617 = stablehlo.broadcast_in_dim %16616, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %16618 = stablehlo.broadcast_in_dim %16617, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16619 = stablehlo.multiply %16618, %4968 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16620 = stablehlo.broadcast_in_dim %16590, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %16621 = stablehlo.broadcast_in_dim %16620, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16622 = stablehlo.multiply %16621, %arg372 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16623 = stablehlo.add %16619, %16622 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %16624 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16625 = stablehlo.broadcast_in_dim %16624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16626 = stablehlo.subtract %16625, %16613 : tensor<1xf32> loc(#loc5260)
    %16627 = stablehlo.multiply %4968, %4968 : tensor<1x3x2048x32x64xf32> loc(#loc5264)
    %16628 = stablehlo.broadcast_in_dim %16626, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %16629 = stablehlo.broadcast_in_dim %16628, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16630 = stablehlo.multiply %16629, %16627 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16631 = stablehlo.broadcast_in_dim %16613, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %16632 = stablehlo.broadcast_in_dim %16631, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16633 = stablehlo.multiply %16632, %arg548 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %16634 = stablehlo.add %16630, %16633 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %16635 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16636 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16637 = stablehlo.broadcast_in_dim %16636, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16638 = stablehlo.add %16635, %16637 : tensor<1xf32> loc(#loc5259)
    %16639 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16640 = stablehlo.broadcast_in_dim %16639, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16641 = stablehlo.subtract %16638, %16640 : tensor<1xf32> loc(#loc5260)
    %16642 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16643 = stablehlo.broadcast_in_dim %16642, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16644 = stablehlo.power %16643, %16641 : tensor<1xf32> loc(#loc5261)
    %16645 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16646 = stablehlo.broadcast_in_dim %16645, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16647 = stablehlo.subtract %16646, %16644 : tensor<1xf32> loc(#loc5260)
    %16648 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16649 = stablehlo.broadcast_in_dim %16648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16650 = stablehlo.multiply %16649, %16647 : tensor<1xf32> loc(#loc5236)
    %16651 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16652 = stablehlo.broadcast_in_dim %16651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16653 = stablehlo.power %16652, %16638 : tensor<1xf32> loc(#loc5261)
    %16654 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16655 = stablehlo.broadcast_in_dim %16654, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16656 = stablehlo.subtract %16655, %16653 : tensor<1xf32> loc(#loc5260)
    %16657 = stablehlo.divide %16650, %16656 : tensor<1xf32> loc(#loc5253)
    %16658 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16659 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16660 = stablehlo.broadcast_in_dim %16659, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16661 = stablehlo.add %16658, %16660 : tensor<1xf32> loc(#loc5259)
    %16662 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16663 = stablehlo.broadcast_in_dim %16662, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16664 = stablehlo.subtract %16661, %16663 : tensor<1xf32> loc(#loc5260)
    %16665 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16666 = stablehlo.broadcast_in_dim %16665, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16667 = stablehlo.power %16666, %16664 : tensor<1xf32> loc(#loc5261)
    %16668 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16669 = stablehlo.broadcast_in_dim %16668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16670 = stablehlo.subtract %16669, %16667 : tensor<1xf32> loc(#loc5260)
    %16671 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16672 = stablehlo.broadcast_in_dim %16671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16673 = stablehlo.multiply %16672, %16670 : tensor<1xf32> loc(#loc5236)
    %16674 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16675 = stablehlo.broadcast_in_dim %16674, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16676 = stablehlo.power %16675, %16661 : tensor<1xf32> loc(#loc5261)
    %16677 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16678 = stablehlo.broadcast_in_dim %16677, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16679 = stablehlo.subtract %16678, %16676 : tensor<1xf32> loc(#loc5260)
    %16680 = stablehlo.divide %16673, %16679 : tensor<1xf32> loc(#loc5253)
    %16681 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16682 = stablehlo.broadcast_in_dim %16681, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16683 = stablehlo.subtract %16682, %16657 : tensor<1xf32> loc(#loc5260)
    %16684 = stablehlo.broadcast_in_dim %16683, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16685 = stablehlo.broadcast_in_dim %16684, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %16686 = stablehlo.multiply %16685, %4972 : tensor<1x64xf32> loc(#loc5262)
    %16687 = stablehlo.broadcast_in_dim %16657, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16688 = stablehlo.broadcast_in_dim %16687, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %16689 = stablehlo.multiply %16688, %arg373 : tensor<1x64xf32> loc(#loc5262)
    %16690 = stablehlo.add %16686, %16689 : tensor<1x64xf32> loc(#loc5263)
    %16691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16692 = stablehlo.broadcast_in_dim %16691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16693 = stablehlo.subtract %16692, %16680 : tensor<1xf32> loc(#loc5260)
    %16694 = stablehlo.multiply %4972, %4972 : tensor<1x64xf32> loc(#loc5264)
    %16695 = stablehlo.broadcast_in_dim %16693, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16696 = stablehlo.broadcast_in_dim %16695, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %16697 = stablehlo.multiply %16696, %16694 : tensor<1x64xf32> loc(#loc5262)
    %16698 = stablehlo.broadcast_in_dim %16680, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %16699 = stablehlo.broadcast_in_dim %16698, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %16700 = stablehlo.multiply %16699, %arg549 : tensor<1x64xf32> loc(#loc5262)
    %16701 = stablehlo.add %16697, %16700 : tensor<1x64xf32> loc(#loc5263)
    %16702 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16703 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16704 = stablehlo.broadcast_in_dim %16703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16705 = stablehlo.add %16702, %16704 : tensor<1xf32> loc(#loc5259)
    %16706 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16707 = stablehlo.broadcast_in_dim %16706, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16708 = stablehlo.subtract %16705, %16707 : tensor<1xf32> loc(#loc5260)
    %16709 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16710 = stablehlo.broadcast_in_dim %16709, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16711 = stablehlo.power %16710, %16708 : tensor<1xf32> loc(#loc5261)
    %16712 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16713 = stablehlo.broadcast_in_dim %16712, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16714 = stablehlo.subtract %16713, %16711 : tensor<1xf32> loc(#loc5260)
    %16715 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16716 = stablehlo.broadcast_in_dim %16715, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16717 = stablehlo.multiply %16716, %16714 : tensor<1xf32> loc(#loc5236)
    %16718 = stablehlo.constant dense<0.899999976> : tensor<f32> loc(#loc)
    %16719 = stablehlo.broadcast_in_dim %16718, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16720 = stablehlo.power %16719, %16705 : tensor<1xf32> loc(#loc5261)
    %16721 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16722 = stablehlo.broadcast_in_dim %16721, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16723 = stablehlo.subtract %16722, %16720 : tensor<1xf32> loc(#loc5260)
    %16724 = stablehlo.divide %16717, %16723 : tensor<1xf32> loc(#loc5253)
    %16725 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5258)
    %16726 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16727 = stablehlo.broadcast_in_dim %16726, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5259)
    %16728 = stablehlo.add %16725, %16727 : tensor<1xf32> loc(#loc5259)
    %16729 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16730 = stablehlo.broadcast_in_dim %16729, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16731 = stablehlo.subtract %16728, %16730 : tensor<1xf32> loc(#loc5260)
    %16732 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16733 = stablehlo.broadcast_in_dim %16732, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16734 = stablehlo.power %16733, %16731 : tensor<1xf32> loc(#loc5261)
    %16735 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16736 = stablehlo.broadcast_in_dim %16735, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16737 = stablehlo.subtract %16736, %16734 : tensor<1xf32> loc(#loc5260)
    %16738 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16739 = stablehlo.broadcast_in_dim %16738, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %16740 = stablehlo.multiply %16739, %16737 : tensor<1xf32> loc(#loc5236)
    %16741 = stablehlo.constant dense<0.949999988> : tensor<f32> loc(#loc)
    %16742 = stablehlo.broadcast_in_dim %16741, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5261)
    %16743 = stablehlo.power %16742, %16728 : tensor<1xf32> loc(#loc5261)
    %16744 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16745 = stablehlo.broadcast_in_dim %16744, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16746 = stablehlo.subtract %16745, %16743 : tensor<1xf32> loc(#loc5260)
    %16747 = stablehlo.divide %16740, %16746 : tensor<1xf32> loc(#loc5253)
    %16748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16749 = stablehlo.broadcast_in_dim %16748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16750 = stablehlo.subtract %16749, %16724 : tensor<1xf32> loc(#loc5260)
    %16751 = stablehlo.broadcast_in_dim %16750, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16752 = stablehlo.broadcast_in_dim %16751, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16753 = stablehlo.multiply %16752, %4976 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16754 = stablehlo.broadcast_in_dim %16724, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16755 = stablehlo.broadcast_in_dim %16754, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16756 = stablehlo.multiply %16755, %arg374 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16757 = stablehlo.add %16753, %16756 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %16758 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %16759 = stablehlo.broadcast_in_dim %16758, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5260)
    %16760 = stablehlo.subtract %16759, %16747 : tensor<1xf32> loc(#loc5260)
    %16761 = stablehlo.multiply %4976, %4976 : tensor<1x2048x32x64xf32> loc(#loc5264)
    %16762 = stablehlo.broadcast_in_dim %16760, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16763 = stablehlo.broadcast_in_dim %16762, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16764 = stablehlo.multiply %16763, %16761 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16765 = stablehlo.broadcast_in_dim %16747, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %16766 = stablehlo.broadcast_in_dim %16765, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %16767 = stablehlo.multiply %16766, %arg550 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %16768 = stablehlo.add %16764, %16767 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %16769 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16770 = stablehlo.broadcast_in_dim %16769, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %16771 = stablehlo.add %5043, %16770 : tensor<1x8192xf32> loc(#loc5265)
    %16772 = stablehlo.sqrt %16771 : tensor<1x8192xf32> loc(#loc5266)
    %16773 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16774 = stablehlo.broadcast_in_dim %16773, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %16775 = stablehlo.add %16772, %16774 : tensor<1x8192xf32> loc(#loc5265)
    %16776 = stablehlo.divide %5032, %16775 : tensor<1x8192xf32> loc(#loc5267)
    %16777 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16778 = stablehlo.broadcast_in_dim %16777, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %16779 = stablehlo.add %5110, %16778 : tensor<1x2048x8192xf32> loc(#loc5265)
    %16780 = stablehlo.sqrt %16779 : tensor<1x2048x8192xf32> loc(#loc5266)
    %16781 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16782 = stablehlo.broadcast_in_dim %16781, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %16783 = stablehlo.add %16780, %16782 : tensor<1x2048x8192xf32> loc(#loc5265)
    %16784 = stablehlo.divide %5099, %16783 : tensor<1x2048x8192xf32> loc(#loc5267)
    %16785 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16786 = stablehlo.broadcast_in_dim %16785, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16787 = stablehlo.add %5177, %16786 : tensor<1x2048xf32> loc(#loc5265)
    %16788 = stablehlo.sqrt %16787 : tensor<1x2048xf32> loc(#loc5266)
    %16789 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16790 = stablehlo.broadcast_in_dim %16789, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16791 = stablehlo.add %16788, %16790 : tensor<1x2048xf32> loc(#loc5265)
    %16792 = stablehlo.divide %5166, %16791 : tensor<1x2048xf32> loc(#loc5267)
    %16793 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16794 = stablehlo.broadcast_in_dim %16793, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %16795 = stablehlo.add %5244, %16794 : tensor<1x8192x2048xf32> loc(#loc5265)
    %16796 = stablehlo.sqrt %16795 : tensor<1x8192x2048xf32> loc(#loc5266)
    %16797 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16798 = stablehlo.broadcast_in_dim %16797, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %16799 = stablehlo.add %16796, %16798 : tensor<1x8192x2048xf32> loc(#loc5265)
    %16800 = stablehlo.divide %5233, %16799 : tensor<1x8192x2048xf32> loc(#loc5267)
    %16801 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16802 = stablehlo.broadcast_in_dim %16801, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16803 = stablehlo.add %5311, %16802 : tensor<1x2048xf32> loc(#loc5265)
    %16804 = stablehlo.sqrt %16803 : tensor<1x2048xf32> loc(#loc5266)
    %16805 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16806 = stablehlo.broadcast_in_dim %16805, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16807 = stablehlo.add %16804, %16806 : tensor<1x2048xf32> loc(#loc5265)
    %16808 = stablehlo.divide %5300, %16807 : tensor<1x2048xf32> loc(#loc5267)
    %16809 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16810 = stablehlo.broadcast_in_dim %16809, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16811 = stablehlo.add %5378, %16810 : tensor<1x2048xf32> loc(#loc5265)
    %16812 = stablehlo.sqrt %16811 : tensor<1x2048xf32> loc(#loc5266)
    %16813 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16814 = stablehlo.broadcast_in_dim %16813, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16815 = stablehlo.add %16812, %16814 : tensor<1x2048xf32> loc(#loc5265)
    %16816 = stablehlo.divide %5367, %16815 : tensor<1x2048xf32> loc(#loc5267)
    %16817 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16818 = stablehlo.broadcast_in_dim %16817, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16819 = stablehlo.add %5445, %16818 : tensor<1x2048xf32> loc(#loc5265)
    %16820 = stablehlo.sqrt %16819 : tensor<1x2048xf32> loc(#loc5266)
    %16821 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16822 = stablehlo.broadcast_in_dim %16821, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16823 = stablehlo.add %16820, %16822 : tensor<1x2048xf32> loc(#loc5265)
    %16824 = stablehlo.divide %5434, %16823 : tensor<1x2048xf32> loc(#loc5267)
    %16825 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16826 = stablehlo.broadcast_in_dim %16825, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16827 = stablehlo.add %5512, %16826 : tensor<1x2048xf32> loc(#loc5265)
    %16828 = stablehlo.sqrt %16827 : tensor<1x2048xf32> loc(#loc5266)
    %16829 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16830 = stablehlo.broadcast_in_dim %16829, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16831 = stablehlo.add %16828, %16830 : tensor<1x2048xf32> loc(#loc5265)
    %16832 = stablehlo.divide %5501, %16831 : tensor<1x2048xf32> loc(#loc5267)
    %16833 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16834 = stablehlo.broadcast_in_dim %16833, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16835 = stablehlo.add %5579, %16834 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16836 = stablehlo.sqrt %16835 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %16837 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16838 = stablehlo.broadcast_in_dim %16837, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16839 = stablehlo.add %16836, %16838 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16840 = stablehlo.divide %5568, %16839 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %16841 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16842 = stablehlo.broadcast_in_dim %16841, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %16843 = stablehlo.add %5646, %16842 : tensor<1x64xf32> loc(#loc5265)
    %16844 = stablehlo.sqrt %16843 : tensor<1x64xf32> loc(#loc5266)
    %16845 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16846 = stablehlo.broadcast_in_dim %16845, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %16847 = stablehlo.add %16844, %16846 : tensor<1x64xf32> loc(#loc5265)
    %16848 = stablehlo.divide %5635, %16847 : tensor<1x64xf32> loc(#loc5267)
    %16849 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16850 = stablehlo.broadcast_in_dim %16849, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %16851 = stablehlo.add %5713, %16850 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %16852 = stablehlo.sqrt %16851 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %16853 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16854 = stablehlo.broadcast_in_dim %16853, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %16855 = stablehlo.add %16852, %16854 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %16856 = stablehlo.divide %5702, %16855 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %16857 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16858 = stablehlo.broadcast_in_dim %16857, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %16859 = stablehlo.add %5780, %16858 : tensor<1x8192xf32> loc(#loc5265)
    %16860 = stablehlo.sqrt %16859 : tensor<1x8192xf32> loc(#loc5266)
    %16861 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16862 = stablehlo.broadcast_in_dim %16861, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %16863 = stablehlo.add %16860, %16862 : tensor<1x8192xf32> loc(#loc5265)
    %16864 = stablehlo.divide %5769, %16863 : tensor<1x8192xf32> loc(#loc5267)
    %16865 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16866 = stablehlo.broadcast_in_dim %16865, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %16867 = stablehlo.add %5847, %16866 : tensor<1x2048x8192xf32> loc(#loc5265)
    %16868 = stablehlo.sqrt %16867 : tensor<1x2048x8192xf32> loc(#loc5266)
    %16869 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16870 = stablehlo.broadcast_in_dim %16869, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %16871 = stablehlo.add %16868, %16870 : tensor<1x2048x8192xf32> loc(#loc5265)
    %16872 = stablehlo.divide %5836, %16871 : tensor<1x2048x8192xf32> loc(#loc5267)
    %16873 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16874 = stablehlo.broadcast_in_dim %16873, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16875 = stablehlo.add %5914, %16874 : tensor<1x2048xf32> loc(#loc5265)
    %16876 = stablehlo.sqrt %16875 : tensor<1x2048xf32> loc(#loc5266)
    %16877 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16878 = stablehlo.broadcast_in_dim %16877, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16879 = stablehlo.add %16876, %16878 : tensor<1x2048xf32> loc(#loc5265)
    %16880 = stablehlo.divide %5903, %16879 : tensor<1x2048xf32> loc(#loc5267)
    %16881 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16882 = stablehlo.broadcast_in_dim %16881, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %16883 = stablehlo.add %5981, %16882 : tensor<1x8192x2048xf32> loc(#loc5265)
    %16884 = stablehlo.sqrt %16883 : tensor<1x8192x2048xf32> loc(#loc5266)
    %16885 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16886 = stablehlo.broadcast_in_dim %16885, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %16887 = stablehlo.add %16884, %16886 : tensor<1x8192x2048xf32> loc(#loc5265)
    %16888 = stablehlo.divide %5970, %16887 : tensor<1x8192x2048xf32> loc(#loc5267)
    %16889 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16890 = stablehlo.broadcast_in_dim %16889, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16891 = stablehlo.add %6048, %16890 : tensor<1x2048xf32> loc(#loc5265)
    %16892 = stablehlo.sqrt %16891 : tensor<1x2048xf32> loc(#loc5266)
    %16893 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16894 = stablehlo.broadcast_in_dim %16893, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16895 = stablehlo.add %16892, %16894 : tensor<1x2048xf32> loc(#loc5265)
    %16896 = stablehlo.divide %6037, %16895 : tensor<1x2048xf32> loc(#loc5267)
    %16897 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16898 = stablehlo.broadcast_in_dim %16897, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16899 = stablehlo.add %6115, %16898 : tensor<1x2048xf32> loc(#loc5265)
    %16900 = stablehlo.sqrt %16899 : tensor<1x2048xf32> loc(#loc5266)
    %16901 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16902 = stablehlo.broadcast_in_dim %16901, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16903 = stablehlo.add %16900, %16902 : tensor<1x2048xf32> loc(#loc5265)
    %16904 = stablehlo.divide %6104, %16903 : tensor<1x2048xf32> loc(#loc5267)
    %16905 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16906 = stablehlo.broadcast_in_dim %16905, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16907 = stablehlo.add %6182, %16906 : tensor<1x2048xf32> loc(#loc5265)
    %16908 = stablehlo.sqrt %16907 : tensor<1x2048xf32> loc(#loc5266)
    %16909 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16910 = stablehlo.broadcast_in_dim %16909, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16911 = stablehlo.add %16908, %16910 : tensor<1x2048xf32> loc(#loc5265)
    %16912 = stablehlo.divide %6171, %16911 : tensor<1x2048xf32> loc(#loc5267)
    %16913 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16914 = stablehlo.broadcast_in_dim %16913, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16915 = stablehlo.add %6249, %16914 : tensor<1x2048xf32> loc(#loc5265)
    %16916 = stablehlo.sqrt %16915 : tensor<1x2048xf32> loc(#loc5266)
    %16917 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16918 = stablehlo.broadcast_in_dim %16917, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16919 = stablehlo.add %16916, %16918 : tensor<1x2048xf32> loc(#loc5265)
    %16920 = stablehlo.divide %6238, %16919 : tensor<1x2048xf32> loc(#loc5267)
    %16921 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16922 = stablehlo.broadcast_in_dim %16921, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16923 = stablehlo.add %6316, %16922 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16924 = stablehlo.sqrt %16923 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %16925 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16926 = stablehlo.broadcast_in_dim %16925, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16927 = stablehlo.add %16924, %16926 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %16928 = stablehlo.divide %6305, %16927 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %16929 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16930 = stablehlo.broadcast_in_dim %16929, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %16931 = stablehlo.add %6383, %16930 : tensor<1x64xf32> loc(#loc5265)
    %16932 = stablehlo.sqrt %16931 : tensor<1x64xf32> loc(#loc5266)
    %16933 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16934 = stablehlo.broadcast_in_dim %16933, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %16935 = stablehlo.add %16932, %16934 : tensor<1x64xf32> loc(#loc5265)
    %16936 = stablehlo.divide %6372, %16935 : tensor<1x64xf32> loc(#loc5267)
    %16937 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16938 = stablehlo.broadcast_in_dim %16937, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %16939 = stablehlo.add %6450, %16938 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %16940 = stablehlo.sqrt %16939 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %16941 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16942 = stablehlo.broadcast_in_dim %16941, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %16943 = stablehlo.add %16940, %16942 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %16944 = stablehlo.divide %6439, %16943 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %16945 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16946 = stablehlo.broadcast_in_dim %16945, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %16947 = stablehlo.add %6517, %16946 : tensor<1x8192xf32> loc(#loc5265)
    %16948 = stablehlo.sqrt %16947 : tensor<1x8192xf32> loc(#loc5266)
    %16949 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16950 = stablehlo.broadcast_in_dim %16949, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %16951 = stablehlo.add %16948, %16950 : tensor<1x8192xf32> loc(#loc5265)
    %16952 = stablehlo.divide %6506, %16951 : tensor<1x8192xf32> loc(#loc5267)
    %16953 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16954 = stablehlo.broadcast_in_dim %16953, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %16955 = stablehlo.add %6584, %16954 : tensor<1x2048x8192xf32> loc(#loc5265)
    %16956 = stablehlo.sqrt %16955 : tensor<1x2048x8192xf32> loc(#loc5266)
    %16957 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16958 = stablehlo.broadcast_in_dim %16957, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %16959 = stablehlo.add %16956, %16958 : tensor<1x2048x8192xf32> loc(#loc5265)
    %16960 = stablehlo.divide %6573, %16959 : tensor<1x2048x8192xf32> loc(#loc5267)
    %16961 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16962 = stablehlo.broadcast_in_dim %16961, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16963 = stablehlo.add %6651, %16962 : tensor<1x2048xf32> loc(#loc5265)
    %16964 = stablehlo.sqrt %16963 : tensor<1x2048xf32> loc(#loc5266)
    %16965 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16966 = stablehlo.broadcast_in_dim %16965, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16967 = stablehlo.add %16964, %16966 : tensor<1x2048xf32> loc(#loc5265)
    %16968 = stablehlo.divide %6640, %16967 : tensor<1x2048xf32> loc(#loc5267)
    %16969 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16970 = stablehlo.broadcast_in_dim %16969, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %16971 = stablehlo.add %6718, %16970 : tensor<1x8192x2048xf32> loc(#loc5265)
    %16972 = stablehlo.sqrt %16971 : tensor<1x8192x2048xf32> loc(#loc5266)
    %16973 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16974 = stablehlo.broadcast_in_dim %16973, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %16975 = stablehlo.add %16972, %16974 : tensor<1x8192x2048xf32> loc(#loc5265)
    %16976 = stablehlo.divide %6707, %16975 : tensor<1x8192x2048xf32> loc(#loc5267)
    %16977 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16978 = stablehlo.broadcast_in_dim %16977, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16979 = stablehlo.add %6785, %16978 : tensor<1x2048xf32> loc(#loc5265)
    %16980 = stablehlo.sqrt %16979 : tensor<1x2048xf32> loc(#loc5266)
    %16981 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16982 = stablehlo.broadcast_in_dim %16981, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16983 = stablehlo.add %16980, %16982 : tensor<1x2048xf32> loc(#loc5265)
    %16984 = stablehlo.divide %6774, %16983 : tensor<1x2048xf32> loc(#loc5267)
    %16985 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16986 = stablehlo.broadcast_in_dim %16985, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16987 = stablehlo.add %6852, %16986 : tensor<1x2048xf32> loc(#loc5265)
    %16988 = stablehlo.sqrt %16987 : tensor<1x2048xf32> loc(#loc5266)
    %16989 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16990 = stablehlo.broadcast_in_dim %16989, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16991 = stablehlo.add %16988, %16990 : tensor<1x2048xf32> loc(#loc5265)
    %16992 = stablehlo.divide %6841, %16991 : tensor<1x2048xf32> loc(#loc5267)
    %16993 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %16994 = stablehlo.broadcast_in_dim %16993, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16995 = stablehlo.add %6919, %16994 : tensor<1x2048xf32> loc(#loc5265)
    %16996 = stablehlo.sqrt %16995 : tensor<1x2048xf32> loc(#loc5266)
    %16997 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %16998 = stablehlo.broadcast_in_dim %16997, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %16999 = stablehlo.add %16996, %16998 : tensor<1x2048xf32> loc(#loc5265)
    %17000 = stablehlo.divide %6908, %16999 : tensor<1x2048xf32> loc(#loc5267)
    %17001 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17002 = stablehlo.broadcast_in_dim %17001, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17003 = stablehlo.add %6986, %17002 : tensor<1x2048xf32> loc(#loc5265)
    %17004 = stablehlo.sqrt %17003 : tensor<1x2048xf32> loc(#loc5266)
    %17005 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17006 = stablehlo.broadcast_in_dim %17005, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17007 = stablehlo.add %17004, %17006 : tensor<1x2048xf32> loc(#loc5265)
    %17008 = stablehlo.divide %6975, %17007 : tensor<1x2048xf32> loc(#loc5267)
    %17009 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17010 = stablehlo.broadcast_in_dim %17009, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17011 = stablehlo.add %7053, %17010 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17012 = stablehlo.sqrt %17011 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17013 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17014 = stablehlo.broadcast_in_dim %17013, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17015 = stablehlo.add %17012, %17014 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17016 = stablehlo.divide %7042, %17015 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17017 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17018 = stablehlo.broadcast_in_dim %17017, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17019 = stablehlo.add %7120, %17018 : tensor<1x64xf32> loc(#loc5265)
    %17020 = stablehlo.sqrt %17019 : tensor<1x64xf32> loc(#loc5266)
    %17021 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17022 = stablehlo.broadcast_in_dim %17021, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17023 = stablehlo.add %17020, %17022 : tensor<1x64xf32> loc(#loc5265)
    %17024 = stablehlo.divide %7109, %17023 : tensor<1x64xf32> loc(#loc5267)
    %17025 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17026 = stablehlo.broadcast_in_dim %17025, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17027 = stablehlo.add %7187, %17026 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17028 = stablehlo.sqrt %17027 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17029 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17030 = stablehlo.broadcast_in_dim %17029, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17031 = stablehlo.add %17028, %17030 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17032 = stablehlo.divide %7176, %17031 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17033 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17034 = stablehlo.broadcast_in_dim %17033, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17035 = stablehlo.add %7254, %17034 : tensor<1x8192xf32> loc(#loc5265)
    %17036 = stablehlo.sqrt %17035 : tensor<1x8192xf32> loc(#loc5266)
    %17037 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17038 = stablehlo.broadcast_in_dim %17037, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17039 = stablehlo.add %17036, %17038 : tensor<1x8192xf32> loc(#loc5265)
    %17040 = stablehlo.divide %7243, %17039 : tensor<1x8192xf32> loc(#loc5267)
    %17041 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17042 = stablehlo.broadcast_in_dim %17041, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17043 = stablehlo.add %7321, %17042 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17044 = stablehlo.sqrt %17043 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17045 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17046 = stablehlo.broadcast_in_dim %17045, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17047 = stablehlo.add %17044, %17046 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17048 = stablehlo.divide %7310, %17047 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17049 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17050 = stablehlo.broadcast_in_dim %17049, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17051 = stablehlo.add %7388, %17050 : tensor<1x2048xf32> loc(#loc5265)
    %17052 = stablehlo.sqrt %17051 : tensor<1x2048xf32> loc(#loc5266)
    %17053 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17054 = stablehlo.broadcast_in_dim %17053, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17055 = stablehlo.add %17052, %17054 : tensor<1x2048xf32> loc(#loc5265)
    %17056 = stablehlo.divide %7377, %17055 : tensor<1x2048xf32> loc(#loc5267)
    %17057 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17058 = stablehlo.broadcast_in_dim %17057, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17059 = stablehlo.add %7455, %17058 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17060 = stablehlo.sqrt %17059 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17061 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17062 = stablehlo.broadcast_in_dim %17061, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17063 = stablehlo.add %17060, %17062 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17064 = stablehlo.divide %7444, %17063 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17065 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17066 = stablehlo.broadcast_in_dim %17065, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17067 = stablehlo.add %7522, %17066 : tensor<1x2048xf32> loc(#loc5265)
    %17068 = stablehlo.sqrt %17067 : tensor<1x2048xf32> loc(#loc5266)
    %17069 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17070 = stablehlo.broadcast_in_dim %17069, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17071 = stablehlo.add %17068, %17070 : tensor<1x2048xf32> loc(#loc5265)
    %17072 = stablehlo.divide %7511, %17071 : tensor<1x2048xf32> loc(#loc5267)
    %17073 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17074 = stablehlo.broadcast_in_dim %17073, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17075 = stablehlo.add %7589, %17074 : tensor<1x2048xf32> loc(#loc5265)
    %17076 = stablehlo.sqrt %17075 : tensor<1x2048xf32> loc(#loc5266)
    %17077 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17078 = stablehlo.broadcast_in_dim %17077, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17079 = stablehlo.add %17076, %17078 : tensor<1x2048xf32> loc(#loc5265)
    %17080 = stablehlo.divide %7578, %17079 : tensor<1x2048xf32> loc(#loc5267)
    %17081 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17082 = stablehlo.broadcast_in_dim %17081, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17083 = stablehlo.add %7656, %17082 : tensor<1x2048xf32> loc(#loc5265)
    %17084 = stablehlo.sqrt %17083 : tensor<1x2048xf32> loc(#loc5266)
    %17085 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17086 = stablehlo.broadcast_in_dim %17085, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17087 = stablehlo.add %17084, %17086 : tensor<1x2048xf32> loc(#loc5265)
    %17088 = stablehlo.divide %7645, %17087 : tensor<1x2048xf32> loc(#loc5267)
    %17089 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17090 = stablehlo.broadcast_in_dim %17089, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17091 = stablehlo.add %7723, %17090 : tensor<1x2048xf32> loc(#loc5265)
    %17092 = stablehlo.sqrt %17091 : tensor<1x2048xf32> loc(#loc5266)
    %17093 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17094 = stablehlo.broadcast_in_dim %17093, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17095 = stablehlo.add %17092, %17094 : tensor<1x2048xf32> loc(#loc5265)
    %17096 = stablehlo.divide %7712, %17095 : tensor<1x2048xf32> loc(#loc5267)
    %17097 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17098 = stablehlo.broadcast_in_dim %17097, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17099 = stablehlo.add %7790, %17098 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17100 = stablehlo.sqrt %17099 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17101 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17102 = stablehlo.broadcast_in_dim %17101, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17103 = stablehlo.add %17100, %17102 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17104 = stablehlo.divide %7779, %17103 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17105 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17106 = stablehlo.broadcast_in_dim %17105, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17107 = stablehlo.add %7857, %17106 : tensor<1x64xf32> loc(#loc5265)
    %17108 = stablehlo.sqrt %17107 : tensor<1x64xf32> loc(#loc5266)
    %17109 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17110 = stablehlo.broadcast_in_dim %17109, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17111 = stablehlo.add %17108, %17110 : tensor<1x64xf32> loc(#loc5265)
    %17112 = stablehlo.divide %7846, %17111 : tensor<1x64xf32> loc(#loc5267)
    %17113 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17114 = stablehlo.broadcast_in_dim %17113, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17115 = stablehlo.add %7924, %17114 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17116 = stablehlo.sqrt %17115 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17117 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17118 = stablehlo.broadcast_in_dim %17117, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17119 = stablehlo.add %17116, %17118 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17120 = stablehlo.divide %7913, %17119 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17121 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17122 = stablehlo.broadcast_in_dim %17121, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17123 = stablehlo.add %7991, %17122 : tensor<1x8192xf32> loc(#loc5265)
    %17124 = stablehlo.sqrt %17123 : tensor<1x8192xf32> loc(#loc5266)
    %17125 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17126 = stablehlo.broadcast_in_dim %17125, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17127 = stablehlo.add %17124, %17126 : tensor<1x8192xf32> loc(#loc5265)
    %17128 = stablehlo.divide %7980, %17127 : tensor<1x8192xf32> loc(#loc5267)
    %17129 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17130 = stablehlo.broadcast_in_dim %17129, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17131 = stablehlo.add %8058, %17130 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17132 = stablehlo.sqrt %17131 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17133 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17134 = stablehlo.broadcast_in_dim %17133, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17135 = stablehlo.add %17132, %17134 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17136 = stablehlo.divide %8047, %17135 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17137 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17138 = stablehlo.broadcast_in_dim %17137, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17139 = stablehlo.add %8125, %17138 : tensor<1x2048xf32> loc(#loc5265)
    %17140 = stablehlo.sqrt %17139 : tensor<1x2048xf32> loc(#loc5266)
    %17141 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17142 = stablehlo.broadcast_in_dim %17141, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17143 = stablehlo.add %17140, %17142 : tensor<1x2048xf32> loc(#loc5265)
    %17144 = stablehlo.divide %8114, %17143 : tensor<1x2048xf32> loc(#loc5267)
    %17145 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17146 = stablehlo.broadcast_in_dim %17145, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17147 = stablehlo.add %8192, %17146 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17148 = stablehlo.sqrt %17147 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17149 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17150 = stablehlo.broadcast_in_dim %17149, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17151 = stablehlo.add %17148, %17150 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17152 = stablehlo.divide %8181, %17151 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17153 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17154 = stablehlo.broadcast_in_dim %17153, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17155 = stablehlo.add %8259, %17154 : tensor<1x2048xf32> loc(#loc5265)
    %17156 = stablehlo.sqrt %17155 : tensor<1x2048xf32> loc(#loc5266)
    %17157 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17158 = stablehlo.broadcast_in_dim %17157, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17159 = stablehlo.add %17156, %17158 : tensor<1x2048xf32> loc(#loc5265)
    %17160 = stablehlo.divide %8248, %17159 : tensor<1x2048xf32> loc(#loc5267)
    %17161 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17162 = stablehlo.broadcast_in_dim %17161, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17163 = stablehlo.add %8326, %17162 : tensor<1x2048xf32> loc(#loc5265)
    %17164 = stablehlo.sqrt %17163 : tensor<1x2048xf32> loc(#loc5266)
    %17165 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17166 = stablehlo.broadcast_in_dim %17165, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17167 = stablehlo.add %17164, %17166 : tensor<1x2048xf32> loc(#loc5265)
    %17168 = stablehlo.divide %8315, %17167 : tensor<1x2048xf32> loc(#loc5267)
    %17169 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17170 = stablehlo.broadcast_in_dim %17169, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17171 = stablehlo.add %8393, %17170 : tensor<1x2048xf32> loc(#loc5265)
    %17172 = stablehlo.sqrt %17171 : tensor<1x2048xf32> loc(#loc5266)
    %17173 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17174 = stablehlo.broadcast_in_dim %17173, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17175 = stablehlo.add %17172, %17174 : tensor<1x2048xf32> loc(#loc5265)
    %17176 = stablehlo.divide %8382, %17175 : tensor<1x2048xf32> loc(#loc5267)
    %17177 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17178 = stablehlo.broadcast_in_dim %17177, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17179 = stablehlo.add %8460, %17178 : tensor<1x2048xf32> loc(#loc5265)
    %17180 = stablehlo.sqrt %17179 : tensor<1x2048xf32> loc(#loc5266)
    %17181 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17182 = stablehlo.broadcast_in_dim %17181, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17183 = stablehlo.add %17180, %17182 : tensor<1x2048xf32> loc(#loc5265)
    %17184 = stablehlo.divide %8449, %17183 : tensor<1x2048xf32> loc(#loc5267)
    %17185 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17186 = stablehlo.broadcast_in_dim %17185, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17187 = stablehlo.add %8527, %17186 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17188 = stablehlo.sqrt %17187 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17189 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17190 = stablehlo.broadcast_in_dim %17189, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17191 = stablehlo.add %17188, %17190 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17192 = stablehlo.divide %8516, %17191 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17193 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17194 = stablehlo.broadcast_in_dim %17193, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17195 = stablehlo.add %8594, %17194 : tensor<1x64xf32> loc(#loc5265)
    %17196 = stablehlo.sqrt %17195 : tensor<1x64xf32> loc(#loc5266)
    %17197 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17198 = stablehlo.broadcast_in_dim %17197, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17199 = stablehlo.add %17196, %17198 : tensor<1x64xf32> loc(#loc5265)
    %17200 = stablehlo.divide %8583, %17199 : tensor<1x64xf32> loc(#loc5267)
    %17201 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17202 = stablehlo.broadcast_in_dim %17201, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17203 = stablehlo.add %8661, %17202 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17204 = stablehlo.sqrt %17203 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17205 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17206 = stablehlo.broadcast_in_dim %17205, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17207 = stablehlo.add %17204, %17206 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17208 = stablehlo.divide %8650, %17207 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17209 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17210 = stablehlo.broadcast_in_dim %17209, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17211 = stablehlo.add %8728, %17210 : tensor<1x8192xf32> loc(#loc5265)
    %17212 = stablehlo.sqrt %17211 : tensor<1x8192xf32> loc(#loc5266)
    %17213 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17214 = stablehlo.broadcast_in_dim %17213, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17215 = stablehlo.add %17212, %17214 : tensor<1x8192xf32> loc(#loc5265)
    %17216 = stablehlo.divide %8717, %17215 : tensor<1x8192xf32> loc(#loc5267)
    %17217 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17218 = stablehlo.broadcast_in_dim %17217, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17219 = stablehlo.add %8795, %17218 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17220 = stablehlo.sqrt %17219 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17221 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17222 = stablehlo.broadcast_in_dim %17221, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17223 = stablehlo.add %17220, %17222 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17224 = stablehlo.divide %8784, %17223 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17225 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17226 = stablehlo.broadcast_in_dim %17225, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17227 = stablehlo.add %8862, %17226 : tensor<1x2048xf32> loc(#loc5265)
    %17228 = stablehlo.sqrt %17227 : tensor<1x2048xf32> loc(#loc5266)
    %17229 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17230 = stablehlo.broadcast_in_dim %17229, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17231 = stablehlo.add %17228, %17230 : tensor<1x2048xf32> loc(#loc5265)
    %17232 = stablehlo.divide %8851, %17231 : tensor<1x2048xf32> loc(#loc5267)
    %17233 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17234 = stablehlo.broadcast_in_dim %17233, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17235 = stablehlo.add %8929, %17234 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17236 = stablehlo.sqrt %17235 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17237 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17238 = stablehlo.broadcast_in_dim %17237, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17239 = stablehlo.add %17236, %17238 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17240 = stablehlo.divide %8918, %17239 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17241 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17242 = stablehlo.broadcast_in_dim %17241, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17243 = stablehlo.add %8996, %17242 : tensor<1x2048xf32> loc(#loc5265)
    %17244 = stablehlo.sqrt %17243 : tensor<1x2048xf32> loc(#loc5266)
    %17245 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17246 = stablehlo.broadcast_in_dim %17245, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17247 = stablehlo.add %17244, %17246 : tensor<1x2048xf32> loc(#loc5265)
    %17248 = stablehlo.divide %8985, %17247 : tensor<1x2048xf32> loc(#loc5267)
    %17249 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17250 = stablehlo.broadcast_in_dim %17249, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17251 = stablehlo.add %9063, %17250 : tensor<1x2048xf32> loc(#loc5265)
    %17252 = stablehlo.sqrt %17251 : tensor<1x2048xf32> loc(#loc5266)
    %17253 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17254 = stablehlo.broadcast_in_dim %17253, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17255 = stablehlo.add %17252, %17254 : tensor<1x2048xf32> loc(#loc5265)
    %17256 = stablehlo.divide %9052, %17255 : tensor<1x2048xf32> loc(#loc5267)
    %17257 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17258 = stablehlo.broadcast_in_dim %17257, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17259 = stablehlo.add %9130, %17258 : tensor<1x2048xf32> loc(#loc5265)
    %17260 = stablehlo.sqrt %17259 : tensor<1x2048xf32> loc(#loc5266)
    %17261 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17262 = stablehlo.broadcast_in_dim %17261, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17263 = stablehlo.add %17260, %17262 : tensor<1x2048xf32> loc(#loc5265)
    %17264 = stablehlo.divide %9119, %17263 : tensor<1x2048xf32> loc(#loc5267)
    %17265 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17266 = stablehlo.broadcast_in_dim %17265, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17267 = stablehlo.add %9197, %17266 : tensor<1x2048xf32> loc(#loc5265)
    %17268 = stablehlo.sqrt %17267 : tensor<1x2048xf32> loc(#loc5266)
    %17269 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17270 = stablehlo.broadcast_in_dim %17269, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17271 = stablehlo.add %17268, %17270 : tensor<1x2048xf32> loc(#loc5265)
    %17272 = stablehlo.divide %9186, %17271 : tensor<1x2048xf32> loc(#loc5267)
    %17273 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17274 = stablehlo.broadcast_in_dim %17273, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17275 = stablehlo.add %9264, %17274 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17276 = stablehlo.sqrt %17275 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17277 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17278 = stablehlo.broadcast_in_dim %17277, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17279 = stablehlo.add %17276, %17278 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17280 = stablehlo.divide %9253, %17279 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17281 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17282 = stablehlo.broadcast_in_dim %17281, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17283 = stablehlo.add %9331, %17282 : tensor<1x64xf32> loc(#loc5265)
    %17284 = stablehlo.sqrt %17283 : tensor<1x64xf32> loc(#loc5266)
    %17285 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17286 = stablehlo.broadcast_in_dim %17285, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17287 = stablehlo.add %17284, %17286 : tensor<1x64xf32> loc(#loc5265)
    %17288 = stablehlo.divide %9320, %17287 : tensor<1x64xf32> loc(#loc5267)
    %17289 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17290 = stablehlo.broadcast_in_dim %17289, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17291 = stablehlo.add %9398, %17290 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17292 = stablehlo.sqrt %17291 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17293 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17294 = stablehlo.broadcast_in_dim %17293, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17295 = stablehlo.add %17292, %17294 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17296 = stablehlo.divide %9387, %17295 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17297 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17298 = stablehlo.broadcast_in_dim %17297, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17299 = stablehlo.add %9465, %17298 : tensor<1x8192xf32> loc(#loc5265)
    %17300 = stablehlo.sqrt %17299 : tensor<1x8192xf32> loc(#loc5266)
    %17301 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17302 = stablehlo.broadcast_in_dim %17301, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17303 = stablehlo.add %17300, %17302 : tensor<1x8192xf32> loc(#loc5265)
    %17304 = stablehlo.divide %9454, %17303 : tensor<1x8192xf32> loc(#loc5267)
    %17305 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17306 = stablehlo.broadcast_in_dim %17305, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17307 = stablehlo.add %9532, %17306 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17308 = stablehlo.sqrt %17307 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17309 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17310 = stablehlo.broadcast_in_dim %17309, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17311 = stablehlo.add %17308, %17310 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17312 = stablehlo.divide %9521, %17311 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17313 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17314 = stablehlo.broadcast_in_dim %17313, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17315 = stablehlo.add %9599, %17314 : tensor<1x2048xf32> loc(#loc5265)
    %17316 = stablehlo.sqrt %17315 : tensor<1x2048xf32> loc(#loc5266)
    %17317 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17318 = stablehlo.broadcast_in_dim %17317, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17319 = stablehlo.add %17316, %17318 : tensor<1x2048xf32> loc(#loc5265)
    %17320 = stablehlo.divide %9588, %17319 : tensor<1x2048xf32> loc(#loc5267)
    %17321 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17322 = stablehlo.broadcast_in_dim %17321, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17323 = stablehlo.add %9666, %17322 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17324 = stablehlo.sqrt %17323 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17325 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17326 = stablehlo.broadcast_in_dim %17325, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17327 = stablehlo.add %17324, %17326 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17328 = stablehlo.divide %9655, %17327 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17329 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17330 = stablehlo.broadcast_in_dim %17329, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17331 = stablehlo.add %9733, %17330 : tensor<1x2048xf32> loc(#loc5265)
    %17332 = stablehlo.sqrt %17331 : tensor<1x2048xf32> loc(#loc5266)
    %17333 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17334 = stablehlo.broadcast_in_dim %17333, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17335 = stablehlo.add %17332, %17334 : tensor<1x2048xf32> loc(#loc5265)
    %17336 = stablehlo.divide %9722, %17335 : tensor<1x2048xf32> loc(#loc5267)
    %17337 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17338 = stablehlo.broadcast_in_dim %17337, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17339 = stablehlo.add %9800, %17338 : tensor<1x2048xf32> loc(#loc5265)
    %17340 = stablehlo.sqrt %17339 : tensor<1x2048xf32> loc(#loc5266)
    %17341 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17342 = stablehlo.broadcast_in_dim %17341, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17343 = stablehlo.add %17340, %17342 : tensor<1x2048xf32> loc(#loc5265)
    %17344 = stablehlo.divide %9789, %17343 : tensor<1x2048xf32> loc(#loc5267)
    %17345 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17346 = stablehlo.broadcast_in_dim %17345, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17347 = stablehlo.add %9867, %17346 : tensor<1x2048xf32> loc(#loc5265)
    %17348 = stablehlo.sqrt %17347 : tensor<1x2048xf32> loc(#loc5266)
    %17349 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17350 = stablehlo.broadcast_in_dim %17349, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17351 = stablehlo.add %17348, %17350 : tensor<1x2048xf32> loc(#loc5265)
    %17352 = stablehlo.divide %9856, %17351 : tensor<1x2048xf32> loc(#loc5267)
    %17353 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17354 = stablehlo.broadcast_in_dim %17353, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17355 = stablehlo.add %9934, %17354 : tensor<1x2048xf32> loc(#loc5265)
    %17356 = stablehlo.sqrt %17355 : tensor<1x2048xf32> loc(#loc5266)
    %17357 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17358 = stablehlo.broadcast_in_dim %17357, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17359 = stablehlo.add %17356, %17358 : tensor<1x2048xf32> loc(#loc5265)
    %17360 = stablehlo.divide %9923, %17359 : tensor<1x2048xf32> loc(#loc5267)
    %17361 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17362 = stablehlo.broadcast_in_dim %17361, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17363 = stablehlo.add %10001, %17362 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17364 = stablehlo.sqrt %17363 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17365 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17366 = stablehlo.broadcast_in_dim %17365, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17367 = stablehlo.add %17364, %17366 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17368 = stablehlo.divide %9990, %17367 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17369 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17370 = stablehlo.broadcast_in_dim %17369, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17371 = stablehlo.add %10068, %17370 : tensor<1x64xf32> loc(#loc5265)
    %17372 = stablehlo.sqrt %17371 : tensor<1x64xf32> loc(#loc5266)
    %17373 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17374 = stablehlo.broadcast_in_dim %17373, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17375 = stablehlo.add %17372, %17374 : tensor<1x64xf32> loc(#loc5265)
    %17376 = stablehlo.divide %10057, %17375 : tensor<1x64xf32> loc(#loc5267)
    %17377 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17378 = stablehlo.broadcast_in_dim %17377, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17379 = stablehlo.add %10135, %17378 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17380 = stablehlo.sqrt %17379 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17381 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17382 = stablehlo.broadcast_in_dim %17381, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17383 = stablehlo.add %17380, %17382 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17384 = stablehlo.divide %10124, %17383 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17385 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17386 = stablehlo.broadcast_in_dim %17385, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17387 = stablehlo.add %10202, %17386 : tensor<1x8192xf32> loc(#loc5265)
    %17388 = stablehlo.sqrt %17387 : tensor<1x8192xf32> loc(#loc5266)
    %17389 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17390 = stablehlo.broadcast_in_dim %17389, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17391 = stablehlo.add %17388, %17390 : tensor<1x8192xf32> loc(#loc5265)
    %17392 = stablehlo.divide %10191, %17391 : tensor<1x8192xf32> loc(#loc5267)
    %17393 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17394 = stablehlo.broadcast_in_dim %17393, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17395 = stablehlo.add %10269, %17394 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17396 = stablehlo.sqrt %17395 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17397 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17398 = stablehlo.broadcast_in_dim %17397, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17399 = stablehlo.add %17396, %17398 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17400 = stablehlo.divide %10258, %17399 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17401 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17402 = stablehlo.broadcast_in_dim %17401, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17403 = stablehlo.add %10336, %17402 : tensor<1x2048xf32> loc(#loc5265)
    %17404 = stablehlo.sqrt %17403 : tensor<1x2048xf32> loc(#loc5266)
    %17405 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17406 = stablehlo.broadcast_in_dim %17405, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17407 = stablehlo.add %17404, %17406 : tensor<1x2048xf32> loc(#loc5265)
    %17408 = stablehlo.divide %10325, %17407 : tensor<1x2048xf32> loc(#loc5267)
    %17409 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17410 = stablehlo.broadcast_in_dim %17409, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17411 = stablehlo.add %10403, %17410 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17412 = stablehlo.sqrt %17411 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17413 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17414 = stablehlo.broadcast_in_dim %17413, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17415 = stablehlo.add %17412, %17414 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17416 = stablehlo.divide %10392, %17415 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17417 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17418 = stablehlo.broadcast_in_dim %17417, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17419 = stablehlo.add %10470, %17418 : tensor<1x2048xf32> loc(#loc5265)
    %17420 = stablehlo.sqrt %17419 : tensor<1x2048xf32> loc(#loc5266)
    %17421 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17422 = stablehlo.broadcast_in_dim %17421, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17423 = stablehlo.add %17420, %17422 : tensor<1x2048xf32> loc(#loc5265)
    %17424 = stablehlo.divide %10459, %17423 : tensor<1x2048xf32> loc(#loc5267)
    %17425 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17426 = stablehlo.broadcast_in_dim %17425, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17427 = stablehlo.add %10537, %17426 : tensor<1x2048xf32> loc(#loc5265)
    %17428 = stablehlo.sqrt %17427 : tensor<1x2048xf32> loc(#loc5266)
    %17429 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17430 = stablehlo.broadcast_in_dim %17429, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17431 = stablehlo.add %17428, %17430 : tensor<1x2048xf32> loc(#loc5265)
    %17432 = stablehlo.divide %10526, %17431 : tensor<1x2048xf32> loc(#loc5267)
    %17433 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17434 = stablehlo.broadcast_in_dim %17433, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17435 = stablehlo.add %10604, %17434 : tensor<1x2048xf32> loc(#loc5265)
    %17436 = stablehlo.sqrt %17435 : tensor<1x2048xf32> loc(#loc5266)
    %17437 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17438 = stablehlo.broadcast_in_dim %17437, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17439 = stablehlo.add %17436, %17438 : tensor<1x2048xf32> loc(#loc5265)
    %17440 = stablehlo.divide %10593, %17439 : tensor<1x2048xf32> loc(#loc5267)
    %17441 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17442 = stablehlo.broadcast_in_dim %17441, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17443 = stablehlo.add %10671, %17442 : tensor<1x2048xf32> loc(#loc5265)
    %17444 = stablehlo.sqrt %17443 : tensor<1x2048xf32> loc(#loc5266)
    %17445 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17446 = stablehlo.broadcast_in_dim %17445, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17447 = stablehlo.add %17444, %17446 : tensor<1x2048xf32> loc(#loc5265)
    %17448 = stablehlo.divide %10660, %17447 : tensor<1x2048xf32> loc(#loc5267)
    %17449 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17450 = stablehlo.broadcast_in_dim %17449, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17451 = stablehlo.add %10738, %17450 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17452 = stablehlo.sqrt %17451 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17453 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17454 = stablehlo.broadcast_in_dim %17453, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17455 = stablehlo.add %17452, %17454 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17456 = stablehlo.divide %10727, %17455 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17457 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17458 = stablehlo.broadcast_in_dim %17457, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17459 = stablehlo.add %10805, %17458 : tensor<1x64xf32> loc(#loc5265)
    %17460 = stablehlo.sqrt %17459 : tensor<1x64xf32> loc(#loc5266)
    %17461 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17462 = stablehlo.broadcast_in_dim %17461, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17463 = stablehlo.add %17460, %17462 : tensor<1x64xf32> loc(#loc5265)
    %17464 = stablehlo.divide %10794, %17463 : tensor<1x64xf32> loc(#loc5267)
    %17465 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17466 = stablehlo.broadcast_in_dim %17465, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17467 = stablehlo.add %10872, %17466 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17468 = stablehlo.sqrt %17467 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17469 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17470 = stablehlo.broadcast_in_dim %17469, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17471 = stablehlo.add %17468, %17470 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17472 = stablehlo.divide %10861, %17471 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17473 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17474 = stablehlo.broadcast_in_dim %17473, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17475 = stablehlo.add %10939, %17474 : tensor<1x8192xf32> loc(#loc5265)
    %17476 = stablehlo.sqrt %17475 : tensor<1x8192xf32> loc(#loc5266)
    %17477 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17478 = stablehlo.broadcast_in_dim %17477, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17479 = stablehlo.add %17476, %17478 : tensor<1x8192xf32> loc(#loc5265)
    %17480 = stablehlo.divide %10928, %17479 : tensor<1x8192xf32> loc(#loc5267)
    %17481 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17482 = stablehlo.broadcast_in_dim %17481, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17483 = stablehlo.add %11006, %17482 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17484 = stablehlo.sqrt %17483 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17485 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17486 = stablehlo.broadcast_in_dim %17485, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17487 = stablehlo.add %17484, %17486 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17488 = stablehlo.divide %10995, %17487 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17489 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17490 = stablehlo.broadcast_in_dim %17489, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17491 = stablehlo.add %11073, %17490 : tensor<1x2048xf32> loc(#loc5265)
    %17492 = stablehlo.sqrt %17491 : tensor<1x2048xf32> loc(#loc5266)
    %17493 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17494 = stablehlo.broadcast_in_dim %17493, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17495 = stablehlo.add %17492, %17494 : tensor<1x2048xf32> loc(#loc5265)
    %17496 = stablehlo.divide %11062, %17495 : tensor<1x2048xf32> loc(#loc5267)
    %17497 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17498 = stablehlo.broadcast_in_dim %17497, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17499 = stablehlo.add %11140, %17498 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17500 = stablehlo.sqrt %17499 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17501 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17502 = stablehlo.broadcast_in_dim %17501, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17503 = stablehlo.add %17500, %17502 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17504 = stablehlo.divide %11129, %17503 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17505 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17506 = stablehlo.broadcast_in_dim %17505, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17507 = stablehlo.add %11207, %17506 : tensor<1x2048xf32> loc(#loc5265)
    %17508 = stablehlo.sqrt %17507 : tensor<1x2048xf32> loc(#loc5266)
    %17509 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17510 = stablehlo.broadcast_in_dim %17509, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17511 = stablehlo.add %17508, %17510 : tensor<1x2048xf32> loc(#loc5265)
    %17512 = stablehlo.divide %11196, %17511 : tensor<1x2048xf32> loc(#loc5267)
    %17513 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17514 = stablehlo.broadcast_in_dim %17513, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17515 = stablehlo.add %11274, %17514 : tensor<1x2048xf32> loc(#loc5265)
    %17516 = stablehlo.sqrt %17515 : tensor<1x2048xf32> loc(#loc5266)
    %17517 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17518 = stablehlo.broadcast_in_dim %17517, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17519 = stablehlo.add %17516, %17518 : tensor<1x2048xf32> loc(#loc5265)
    %17520 = stablehlo.divide %11263, %17519 : tensor<1x2048xf32> loc(#loc5267)
    %17521 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17522 = stablehlo.broadcast_in_dim %17521, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17523 = stablehlo.add %11341, %17522 : tensor<1x2048xf32> loc(#loc5265)
    %17524 = stablehlo.sqrt %17523 : tensor<1x2048xf32> loc(#loc5266)
    %17525 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17526 = stablehlo.broadcast_in_dim %17525, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17527 = stablehlo.add %17524, %17526 : tensor<1x2048xf32> loc(#loc5265)
    %17528 = stablehlo.divide %11330, %17527 : tensor<1x2048xf32> loc(#loc5267)
    %17529 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17530 = stablehlo.broadcast_in_dim %17529, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17531 = stablehlo.add %11408, %17530 : tensor<1x2048xf32> loc(#loc5265)
    %17532 = stablehlo.sqrt %17531 : tensor<1x2048xf32> loc(#loc5266)
    %17533 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17534 = stablehlo.broadcast_in_dim %17533, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17535 = stablehlo.add %17532, %17534 : tensor<1x2048xf32> loc(#loc5265)
    %17536 = stablehlo.divide %11397, %17535 : tensor<1x2048xf32> loc(#loc5267)
    %17537 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17538 = stablehlo.broadcast_in_dim %17537, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17539 = stablehlo.add %11475, %17538 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17540 = stablehlo.sqrt %17539 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17541 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17542 = stablehlo.broadcast_in_dim %17541, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17543 = stablehlo.add %17540, %17542 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17544 = stablehlo.divide %11464, %17543 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17545 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17546 = stablehlo.broadcast_in_dim %17545, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17547 = stablehlo.add %11542, %17546 : tensor<1x64xf32> loc(#loc5265)
    %17548 = stablehlo.sqrt %17547 : tensor<1x64xf32> loc(#loc5266)
    %17549 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17550 = stablehlo.broadcast_in_dim %17549, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17551 = stablehlo.add %17548, %17550 : tensor<1x64xf32> loc(#loc5265)
    %17552 = stablehlo.divide %11531, %17551 : tensor<1x64xf32> loc(#loc5267)
    %17553 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17554 = stablehlo.broadcast_in_dim %17553, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17555 = stablehlo.add %11609, %17554 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17556 = stablehlo.sqrt %17555 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17557 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17558 = stablehlo.broadcast_in_dim %17557, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17559 = stablehlo.add %17556, %17558 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17560 = stablehlo.divide %11598, %17559 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17561 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17562 = stablehlo.broadcast_in_dim %17561, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17563 = stablehlo.add %11676, %17562 : tensor<1x8192xf32> loc(#loc5265)
    %17564 = stablehlo.sqrt %17563 : tensor<1x8192xf32> loc(#loc5266)
    %17565 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17566 = stablehlo.broadcast_in_dim %17565, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17567 = stablehlo.add %17564, %17566 : tensor<1x8192xf32> loc(#loc5265)
    %17568 = stablehlo.divide %11665, %17567 : tensor<1x8192xf32> loc(#loc5267)
    %17569 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17570 = stablehlo.broadcast_in_dim %17569, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17571 = stablehlo.add %11743, %17570 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17572 = stablehlo.sqrt %17571 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17573 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17574 = stablehlo.broadcast_in_dim %17573, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17575 = stablehlo.add %17572, %17574 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17576 = stablehlo.divide %11732, %17575 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17577 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17578 = stablehlo.broadcast_in_dim %17577, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17579 = stablehlo.add %11810, %17578 : tensor<1x2048xf32> loc(#loc5265)
    %17580 = stablehlo.sqrt %17579 : tensor<1x2048xf32> loc(#loc5266)
    %17581 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17582 = stablehlo.broadcast_in_dim %17581, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17583 = stablehlo.add %17580, %17582 : tensor<1x2048xf32> loc(#loc5265)
    %17584 = stablehlo.divide %11799, %17583 : tensor<1x2048xf32> loc(#loc5267)
    %17585 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17586 = stablehlo.broadcast_in_dim %17585, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17587 = stablehlo.add %11877, %17586 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17588 = stablehlo.sqrt %17587 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17589 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17590 = stablehlo.broadcast_in_dim %17589, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17591 = stablehlo.add %17588, %17590 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17592 = stablehlo.divide %11866, %17591 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17593 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17594 = stablehlo.broadcast_in_dim %17593, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17595 = stablehlo.add %11944, %17594 : tensor<1x2048xf32> loc(#loc5265)
    %17596 = stablehlo.sqrt %17595 : tensor<1x2048xf32> loc(#loc5266)
    %17597 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17598 = stablehlo.broadcast_in_dim %17597, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17599 = stablehlo.add %17596, %17598 : tensor<1x2048xf32> loc(#loc5265)
    %17600 = stablehlo.divide %11933, %17599 : tensor<1x2048xf32> loc(#loc5267)
    %17601 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17602 = stablehlo.broadcast_in_dim %17601, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17603 = stablehlo.add %12011, %17602 : tensor<1x2048xf32> loc(#loc5265)
    %17604 = stablehlo.sqrt %17603 : tensor<1x2048xf32> loc(#loc5266)
    %17605 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17606 = stablehlo.broadcast_in_dim %17605, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17607 = stablehlo.add %17604, %17606 : tensor<1x2048xf32> loc(#loc5265)
    %17608 = stablehlo.divide %12000, %17607 : tensor<1x2048xf32> loc(#loc5267)
    %17609 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17610 = stablehlo.broadcast_in_dim %17609, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17611 = stablehlo.add %12078, %17610 : tensor<1x2048xf32> loc(#loc5265)
    %17612 = stablehlo.sqrt %17611 : tensor<1x2048xf32> loc(#loc5266)
    %17613 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17614 = stablehlo.broadcast_in_dim %17613, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17615 = stablehlo.add %17612, %17614 : tensor<1x2048xf32> loc(#loc5265)
    %17616 = stablehlo.divide %12067, %17615 : tensor<1x2048xf32> loc(#loc5267)
    %17617 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17618 = stablehlo.broadcast_in_dim %17617, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17619 = stablehlo.add %12145, %17618 : tensor<1x2048xf32> loc(#loc5265)
    %17620 = stablehlo.sqrt %17619 : tensor<1x2048xf32> loc(#loc5266)
    %17621 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17622 = stablehlo.broadcast_in_dim %17621, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17623 = stablehlo.add %17620, %17622 : tensor<1x2048xf32> loc(#loc5265)
    %17624 = stablehlo.divide %12134, %17623 : tensor<1x2048xf32> loc(#loc5267)
    %17625 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17626 = stablehlo.broadcast_in_dim %17625, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17627 = stablehlo.add %12212, %17626 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17628 = stablehlo.sqrt %17627 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17629 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17630 = stablehlo.broadcast_in_dim %17629, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17631 = stablehlo.add %17628, %17630 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17632 = stablehlo.divide %12201, %17631 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17633 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17634 = stablehlo.broadcast_in_dim %17633, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17635 = stablehlo.add %12279, %17634 : tensor<1x64xf32> loc(#loc5265)
    %17636 = stablehlo.sqrt %17635 : tensor<1x64xf32> loc(#loc5266)
    %17637 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17638 = stablehlo.broadcast_in_dim %17637, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17639 = stablehlo.add %17636, %17638 : tensor<1x64xf32> loc(#loc5265)
    %17640 = stablehlo.divide %12268, %17639 : tensor<1x64xf32> loc(#loc5267)
    %17641 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17642 = stablehlo.broadcast_in_dim %17641, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17643 = stablehlo.add %12346, %17642 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17644 = stablehlo.sqrt %17643 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17645 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17646 = stablehlo.broadcast_in_dim %17645, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17647 = stablehlo.add %17644, %17646 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17648 = stablehlo.divide %12335, %17647 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17649 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17650 = stablehlo.broadcast_in_dim %17649, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17651 = stablehlo.add %12413, %17650 : tensor<1x8192xf32> loc(#loc5265)
    %17652 = stablehlo.sqrt %17651 : tensor<1x8192xf32> loc(#loc5266)
    %17653 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17654 = stablehlo.broadcast_in_dim %17653, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17655 = stablehlo.add %17652, %17654 : tensor<1x8192xf32> loc(#loc5265)
    %17656 = stablehlo.divide %12402, %17655 : tensor<1x8192xf32> loc(#loc5267)
    %17657 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17658 = stablehlo.broadcast_in_dim %17657, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17659 = stablehlo.add %12480, %17658 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17660 = stablehlo.sqrt %17659 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17661 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17662 = stablehlo.broadcast_in_dim %17661, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17663 = stablehlo.add %17660, %17662 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17664 = stablehlo.divide %12469, %17663 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17665 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17666 = stablehlo.broadcast_in_dim %17665, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17667 = stablehlo.add %12547, %17666 : tensor<1x2048xf32> loc(#loc5265)
    %17668 = stablehlo.sqrt %17667 : tensor<1x2048xf32> loc(#loc5266)
    %17669 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17670 = stablehlo.broadcast_in_dim %17669, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17671 = stablehlo.add %17668, %17670 : tensor<1x2048xf32> loc(#loc5265)
    %17672 = stablehlo.divide %12536, %17671 : tensor<1x2048xf32> loc(#loc5267)
    %17673 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17674 = stablehlo.broadcast_in_dim %17673, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17675 = stablehlo.add %12614, %17674 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17676 = stablehlo.sqrt %17675 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17677 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17678 = stablehlo.broadcast_in_dim %17677, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17679 = stablehlo.add %17676, %17678 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17680 = stablehlo.divide %12603, %17679 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17681 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17682 = stablehlo.broadcast_in_dim %17681, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17683 = stablehlo.add %12681, %17682 : tensor<1x2048xf32> loc(#loc5265)
    %17684 = stablehlo.sqrt %17683 : tensor<1x2048xf32> loc(#loc5266)
    %17685 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17686 = stablehlo.broadcast_in_dim %17685, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17687 = stablehlo.add %17684, %17686 : tensor<1x2048xf32> loc(#loc5265)
    %17688 = stablehlo.divide %12670, %17687 : tensor<1x2048xf32> loc(#loc5267)
    %17689 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17690 = stablehlo.broadcast_in_dim %17689, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17691 = stablehlo.add %12748, %17690 : tensor<1x2048xf32> loc(#loc5265)
    %17692 = stablehlo.sqrt %17691 : tensor<1x2048xf32> loc(#loc5266)
    %17693 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17694 = stablehlo.broadcast_in_dim %17693, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17695 = stablehlo.add %17692, %17694 : tensor<1x2048xf32> loc(#loc5265)
    %17696 = stablehlo.divide %12737, %17695 : tensor<1x2048xf32> loc(#loc5267)
    %17697 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17698 = stablehlo.broadcast_in_dim %17697, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17699 = stablehlo.add %12815, %17698 : tensor<1x2048xf32> loc(#loc5265)
    %17700 = stablehlo.sqrt %17699 : tensor<1x2048xf32> loc(#loc5266)
    %17701 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17702 = stablehlo.broadcast_in_dim %17701, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17703 = stablehlo.add %17700, %17702 : tensor<1x2048xf32> loc(#loc5265)
    %17704 = stablehlo.divide %12804, %17703 : tensor<1x2048xf32> loc(#loc5267)
    %17705 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17706 = stablehlo.broadcast_in_dim %17705, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17707 = stablehlo.add %12882, %17706 : tensor<1x2048xf32> loc(#loc5265)
    %17708 = stablehlo.sqrt %17707 : tensor<1x2048xf32> loc(#loc5266)
    %17709 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17710 = stablehlo.broadcast_in_dim %17709, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17711 = stablehlo.add %17708, %17710 : tensor<1x2048xf32> loc(#loc5265)
    %17712 = stablehlo.divide %12871, %17711 : tensor<1x2048xf32> loc(#loc5267)
    %17713 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17714 = stablehlo.broadcast_in_dim %17713, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17715 = stablehlo.add %12949, %17714 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17716 = stablehlo.sqrt %17715 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17717 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17718 = stablehlo.broadcast_in_dim %17717, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17719 = stablehlo.add %17716, %17718 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17720 = stablehlo.divide %12938, %17719 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17721 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17722 = stablehlo.broadcast_in_dim %17721, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17723 = stablehlo.add %13016, %17722 : tensor<1x64xf32> loc(#loc5265)
    %17724 = stablehlo.sqrt %17723 : tensor<1x64xf32> loc(#loc5266)
    %17725 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17726 = stablehlo.broadcast_in_dim %17725, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17727 = stablehlo.add %17724, %17726 : tensor<1x64xf32> loc(#loc5265)
    %17728 = stablehlo.divide %13005, %17727 : tensor<1x64xf32> loc(#loc5267)
    %17729 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17730 = stablehlo.broadcast_in_dim %17729, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17731 = stablehlo.add %13083, %17730 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17732 = stablehlo.sqrt %17731 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17733 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17734 = stablehlo.broadcast_in_dim %17733, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17735 = stablehlo.add %17732, %17734 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17736 = stablehlo.divide %13072, %17735 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17737 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17738 = stablehlo.broadcast_in_dim %17737, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17739 = stablehlo.add %13150, %17738 : tensor<1x8192xf32> loc(#loc5265)
    %17740 = stablehlo.sqrt %17739 : tensor<1x8192xf32> loc(#loc5266)
    %17741 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17742 = stablehlo.broadcast_in_dim %17741, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17743 = stablehlo.add %17740, %17742 : tensor<1x8192xf32> loc(#loc5265)
    %17744 = stablehlo.divide %13139, %17743 : tensor<1x8192xf32> loc(#loc5267)
    %17745 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17746 = stablehlo.broadcast_in_dim %17745, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17747 = stablehlo.add %13217, %17746 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17748 = stablehlo.sqrt %17747 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17749 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17750 = stablehlo.broadcast_in_dim %17749, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17751 = stablehlo.add %17748, %17750 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17752 = stablehlo.divide %13206, %17751 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17753 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17754 = stablehlo.broadcast_in_dim %17753, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17755 = stablehlo.add %13284, %17754 : tensor<1x2048xf32> loc(#loc5265)
    %17756 = stablehlo.sqrt %17755 : tensor<1x2048xf32> loc(#loc5266)
    %17757 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17758 = stablehlo.broadcast_in_dim %17757, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17759 = stablehlo.add %17756, %17758 : tensor<1x2048xf32> loc(#loc5265)
    %17760 = stablehlo.divide %13273, %17759 : tensor<1x2048xf32> loc(#loc5267)
    %17761 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17762 = stablehlo.broadcast_in_dim %17761, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17763 = stablehlo.add %13351, %17762 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17764 = stablehlo.sqrt %17763 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17765 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17766 = stablehlo.broadcast_in_dim %17765, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17767 = stablehlo.add %17764, %17766 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17768 = stablehlo.divide %13340, %17767 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17769 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17770 = stablehlo.broadcast_in_dim %17769, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17771 = stablehlo.add %13418, %17770 : tensor<1x2048xf32> loc(#loc5265)
    %17772 = stablehlo.sqrt %17771 : tensor<1x2048xf32> loc(#loc5266)
    %17773 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17774 = stablehlo.broadcast_in_dim %17773, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17775 = stablehlo.add %17772, %17774 : tensor<1x2048xf32> loc(#loc5265)
    %17776 = stablehlo.divide %13407, %17775 : tensor<1x2048xf32> loc(#loc5267)
    %17777 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17778 = stablehlo.broadcast_in_dim %17777, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17779 = stablehlo.add %13485, %17778 : tensor<1x2048xf32> loc(#loc5265)
    %17780 = stablehlo.sqrt %17779 : tensor<1x2048xf32> loc(#loc5266)
    %17781 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17782 = stablehlo.broadcast_in_dim %17781, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17783 = stablehlo.add %17780, %17782 : tensor<1x2048xf32> loc(#loc5265)
    %17784 = stablehlo.divide %13474, %17783 : tensor<1x2048xf32> loc(#loc5267)
    %17785 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17786 = stablehlo.broadcast_in_dim %17785, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17787 = stablehlo.add %13552, %17786 : tensor<1x2048xf32> loc(#loc5265)
    %17788 = stablehlo.sqrt %17787 : tensor<1x2048xf32> loc(#loc5266)
    %17789 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17790 = stablehlo.broadcast_in_dim %17789, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17791 = stablehlo.add %17788, %17790 : tensor<1x2048xf32> loc(#loc5265)
    %17792 = stablehlo.divide %13541, %17791 : tensor<1x2048xf32> loc(#loc5267)
    %17793 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17794 = stablehlo.broadcast_in_dim %17793, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17795 = stablehlo.add %13619, %17794 : tensor<1x2048xf32> loc(#loc5265)
    %17796 = stablehlo.sqrt %17795 : tensor<1x2048xf32> loc(#loc5266)
    %17797 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17798 = stablehlo.broadcast_in_dim %17797, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17799 = stablehlo.add %17796, %17798 : tensor<1x2048xf32> loc(#loc5265)
    %17800 = stablehlo.divide %13608, %17799 : tensor<1x2048xf32> loc(#loc5267)
    %17801 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17802 = stablehlo.broadcast_in_dim %17801, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17803 = stablehlo.add %13686, %17802 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17804 = stablehlo.sqrt %17803 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17805 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17806 = stablehlo.broadcast_in_dim %17805, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17807 = stablehlo.add %17804, %17806 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17808 = stablehlo.divide %13675, %17807 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17809 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17810 = stablehlo.broadcast_in_dim %17809, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17811 = stablehlo.add %13753, %17810 : tensor<1x64xf32> loc(#loc5265)
    %17812 = stablehlo.sqrt %17811 : tensor<1x64xf32> loc(#loc5266)
    %17813 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17814 = stablehlo.broadcast_in_dim %17813, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17815 = stablehlo.add %17812, %17814 : tensor<1x64xf32> loc(#loc5265)
    %17816 = stablehlo.divide %13742, %17815 : tensor<1x64xf32> loc(#loc5267)
    %17817 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17818 = stablehlo.broadcast_in_dim %17817, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17819 = stablehlo.add %13820, %17818 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17820 = stablehlo.sqrt %17819 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17821 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17822 = stablehlo.broadcast_in_dim %17821, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17823 = stablehlo.add %17820, %17822 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17824 = stablehlo.divide %13809, %17823 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17825 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17826 = stablehlo.broadcast_in_dim %17825, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17827 = stablehlo.add %13887, %17826 : tensor<1x8192xf32> loc(#loc5265)
    %17828 = stablehlo.sqrt %17827 : tensor<1x8192xf32> loc(#loc5266)
    %17829 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17830 = stablehlo.broadcast_in_dim %17829, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17831 = stablehlo.add %17828, %17830 : tensor<1x8192xf32> loc(#loc5265)
    %17832 = stablehlo.divide %13876, %17831 : tensor<1x8192xf32> loc(#loc5267)
    %17833 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17834 = stablehlo.broadcast_in_dim %17833, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17835 = stablehlo.add %13954, %17834 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17836 = stablehlo.sqrt %17835 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17837 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17838 = stablehlo.broadcast_in_dim %17837, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17839 = stablehlo.add %17836, %17838 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17840 = stablehlo.divide %13943, %17839 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17841 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17842 = stablehlo.broadcast_in_dim %17841, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17843 = stablehlo.add %14021, %17842 : tensor<1x2048xf32> loc(#loc5265)
    %17844 = stablehlo.sqrt %17843 : tensor<1x2048xf32> loc(#loc5266)
    %17845 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17846 = stablehlo.broadcast_in_dim %17845, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17847 = stablehlo.add %17844, %17846 : tensor<1x2048xf32> loc(#loc5265)
    %17848 = stablehlo.divide %14010, %17847 : tensor<1x2048xf32> loc(#loc5267)
    %17849 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17850 = stablehlo.broadcast_in_dim %17849, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17851 = stablehlo.add %14088, %17850 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17852 = stablehlo.sqrt %17851 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17853 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17854 = stablehlo.broadcast_in_dim %17853, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17855 = stablehlo.add %17852, %17854 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17856 = stablehlo.divide %14077, %17855 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17857 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17858 = stablehlo.broadcast_in_dim %17857, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17859 = stablehlo.add %14155, %17858 : tensor<1x2048xf32> loc(#loc5265)
    %17860 = stablehlo.sqrt %17859 : tensor<1x2048xf32> loc(#loc5266)
    %17861 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17862 = stablehlo.broadcast_in_dim %17861, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17863 = stablehlo.add %17860, %17862 : tensor<1x2048xf32> loc(#loc5265)
    %17864 = stablehlo.divide %14144, %17863 : tensor<1x2048xf32> loc(#loc5267)
    %17865 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17866 = stablehlo.broadcast_in_dim %17865, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17867 = stablehlo.add %14222, %17866 : tensor<1x2048xf32> loc(#loc5265)
    %17868 = stablehlo.sqrt %17867 : tensor<1x2048xf32> loc(#loc5266)
    %17869 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17870 = stablehlo.broadcast_in_dim %17869, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17871 = stablehlo.add %17868, %17870 : tensor<1x2048xf32> loc(#loc5265)
    %17872 = stablehlo.divide %14211, %17871 : tensor<1x2048xf32> loc(#loc5267)
    %17873 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17874 = stablehlo.broadcast_in_dim %17873, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17875 = stablehlo.add %14289, %17874 : tensor<1x2048xf32> loc(#loc5265)
    %17876 = stablehlo.sqrt %17875 : tensor<1x2048xf32> loc(#loc5266)
    %17877 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17878 = stablehlo.broadcast_in_dim %17877, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17879 = stablehlo.add %17876, %17878 : tensor<1x2048xf32> loc(#loc5265)
    %17880 = stablehlo.divide %14278, %17879 : tensor<1x2048xf32> loc(#loc5267)
    %17881 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17882 = stablehlo.broadcast_in_dim %17881, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17883 = stablehlo.add %14356, %17882 : tensor<1x2048xf32> loc(#loc5265)
    %17884 = stablehlo.sqrt %17883 : tensor<1x2048xf32> loc(#loc5266)
    %17885 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17886 = stablehlo.broadcast_in_dim %17885, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17887 = stablehlo.add %17884, %17886 : tensor<1x2048xf32> loc(#loc5265)
    %17888 = stablehlo.divide %14345, %17887 : tensor<1x2048xf32> loc(#loc5267)
    %17889 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17890 = stablehlo.broadcast_in_dim %17889, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17891 = stablehlo.add %14423, %17890 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17892 = stablehlo.sqrt %17891 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17893 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17894 = stablehlo.broadcast_in_dim %17893, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17895 = stablehlo.add %17892, %17894 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17896 = stablehlo.divide %14412, %17895 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17897 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17898 = stablehlo.broadcast_in_dim %17897, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17899 = stablehlo.add %14490, %17898 : tensor<1x64xf32> loc(#loc5265)
    %17900 = stablehlo.sqrt %17899 : tensor<1x64xf32> loc(#loc5266)
    %17901 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17902 = stablehlo.broadcast_in_dim %17901, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17903 = stablehlo.add %17900, %17902 : tensor<1x64xf32> loc(#loc5265)
    %17904 = stablehlo.divide %14479, %17903 : tensor<1x64xf32> loc(#loc5267)
    %17905 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17906 = stablehlo.broadcast_in_dim %17905, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17907 = stablehlo.add %14557, %17906 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17908 = stablehlo.sqrt %17907 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17909 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17910 = stablehlo.broadcast_in_dim %17909, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17911 = stablehlo.add %17908, %17910 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17912 = stablehlo.divide %14546, %17911 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %17913 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17914 = stablehlo.broadcast_in_dim %17913, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17915 = stablehlo.add %14624, %17914 : tensor<1x8192xf32> loc(#loc5265)
    %17916 = stablehlo.sqrt %17915 : tensor<1x8192xf32> loc(#loc5266)
    %17917 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17918 = stablehlo.broadcast_in_dim %17917, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %17919 = stablehlo.add %17916, %17918 : tensor<1x8192xf32> loc(#loc5265)
    %17920 = stablehlo.divide %14613, %17919 : tensor<1x8192xf32> loc(#loc5267)
    %17921 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17922 = stablehlo.broadcast_in_dim %17921, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17923 = stablehlo.add %14691, %17922 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17924 = stablehlo.sqrt %17923 : tensor<1x2048x8192xf32> loc(#loc5266)
    %17925 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17926 = stablehlo.broadcast_in_dim %17925, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %17927 = stablehlo.add %17924, %17926 : tensor<1x2048x8192xf32> loc(#loc5265)
    %17928 = stablehlo.divide %14680, %17927 : tensor<1x2048x8192xf32> loc(#loc5267)
    %17929 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17930 = stablehlo.broadcast_in_dim %17929, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17931 = stablehlo.add %14758, %17930 : tensor<1x2048xf32> loc(#loc5265)
    %17932 = stablehlo.sqrt %17931 : tensor<1x2048xf32> loc(#loc5266)
    %17933 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17934 = stablehlo.broadcast_in_dim %17933, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17935 = stablehlo.add %17932, %17934 : tensor<1x2048xf32> loc(#loc5265)
    %17936 = stablehlo.divide %14747, %17935 : tensor<1x2048xf32> loc(#loc5267)
    %17937 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17938 = stablehlo.broadcast_in_dim %17937, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17939 = stablehlo.add %14825, %17938 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17940 = stablehlo.sqrt %17939 : tensor<1x8192x2048xf32> loc(#loc5266)
    %17941 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17942 = stablehlo.broadcast_in_dim %17941, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %17943 = stablehlo.add %17940, %17942 : tensor<1x8192x2048xf32> loc(#loc5265)
    %17944 = stablehlo.divide %14814, %17943 : tensor<1x8192x2048xf32> loc(#loc5267)
    %17945 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17946 = stablehlo.broadcast_in_dim %17945, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17947 = stablehlo.add %14892, %17946 : tensor<1x2048xf32> loc(#loc5265)
    %17948 = stablehlo.sqrt %17947 : tensor<1x2048xf32> loc(#loc5266)
    %17949 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17950 = stablehlo.broadcast_in_dim %17949, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17951 = stablehlo.add %17948, %17950 : tensor<1x2048xf32> loc(#loc5265)
    %17952 = stablehlo.divide %14881, %17951 : tensor<1x2048xf32> loc(#loc5267)
    %17953 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17954 = stablehlo.broadcast_in_dim %17953, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17955 = stablehlo.add %14959, %17954 : tensor<1x2048xf32> loc(#loc5265)
    %17956 = stablehlo.sqrt %17955 : tensor<1x2048xf32> loc(#loc5266)
    %17957 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17958 = stablehlo.broadcast_in_dim %17957, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17959 = stablehlo.add %17956, %17958 : tensor<1x2048xf32> loc(#loc5265)
    %17960 = stablehlo.divide %14948, %17959 : tensor<1x2048xf32> loc(#loc5267)
    %17961 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17962 = stablehlo.broadcast_in_dim %17961, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17963 = stablehlo.add %15026, %17962 : tensor<1x2048xf32> loc(#loc5265)
    %17964 = stablehlo.sqrt %17963 : tensor<1x2048xf32> loc(#loc5266)
    %17965 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17966 = stablehlo.broadcast_in_dim %17965, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17967 = stablehlo.add %17964, %17966 : tensor<1x2048xf32> loc(#loc5265)
    %17968 = stablehlo.divide %15015, %17967 : tensor<1x2048xf32> loc(#loc5267)
    %17969 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17970 = stablehlo.broadcast_in_dim %17969, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17971 = stablehlo.add %15093, %17970 : tensor<1x2048xf32> loc(#loc5265)
    %17972 = stablehlo.sqrt %17971 : tensor<1x2048xf32> loc(#loc5266)
    %17973 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17974 = stablehlo.broadcast_in_dim %17973, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %17975 = stablehlo.add %17972, %17974 : tensor<1x2048xf32> loc(#loc5265)
    %17976 = stablehlo.divide %15082, %17975 : tensor<1x2048xf32> loc(#loc5267)
    %17977 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17978 = stablehlo.broadcast_in_dim %17977, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17979 = stablehlo.add %15160, %17978 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17980 = stablehlo.sqrt %17979 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %17981 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17982 = stablehlo.broadcast_in_dim %17981, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17983 = stablehlo.add %17980, %17982 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %17984 = stablehlo.divide %15149, %17983 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %17985 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17986 = stablehlo.broadcast_in_dim %17985, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17987 = stablehlo.add %15227, %17986 : tensor<1x64xf32> loc(#loc5265)
    %17988 = stablehlo.sqrt %17987 : tensor<1x64xf32> loc(#loc5266)
    %17989 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17990 = stablehlo.broadcast_in_dim %17989, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %17991 = stablehlo.add %17988, %17990 : tensor<1x64xf32> loc(#loc5265)
    %17992 = stablehlo.divide %15216, %17991 : tensor<1x64xf32> loc(#loc5267)
    %17993 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %17994 = stablehlo.broadcast_in_dim %17993, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17995 = stablehlo.add %15294, %17994 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %17996 = stablehlo.sqrt %17995 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %17997 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %17998 = stablehlo.broadcast_in_dim %17997, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %17999 = stablehlo.add %17996, %17998 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %18000 = stablehlo.divide %15283, %17999 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %18001 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18002 = stablehlo.broadcast_in_dim %18001, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %18003 = stablehlo.add %15361, %18002 : tensor<1x8192xf32> loc(#loc5265)
    %18004 = stablehlo.sqrt %18003 : tensor<1x8192xf32> loc(#loc5266)
    %18005 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18006 = stablehlo.broadcast_in_dim %18005, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %18007 = stablehlo.add %18004, %18006 : tensor<1x8192xf32> loc(#loc5265)
    %18008 = stablehlo.divide %15350, %18007 : tensor<1x8192xf32> loc(#loc5267)
    %18009 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18010 = stablehlo.broadcast_in_dim %18009, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %18011 = stablehlo.add %15428, %18010 : tensor<1x2048x8192xf32> loc(#loc5265)
    %18012 = stablehlo.sqrt %18011 : tensor<1x2048x8192xf32> loc(#loc5266)
    %18013 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18014 = stablehlo.broadcast_in_dim %18013, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %18015 = stablehlo.add %18012, %18014 : tensor<1x2048x8192xf32> loc(#loc5265)
    %18016 = stablehlo.divide %15417, %18015 : tensor<1x2048x8192xf32> loc(#loc5267)
    %18017 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18018 = stablehlo.broadcast_in_dim %18017, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18019 = stablehlo.add %15495, %18018 : tensor<1x2048xf32> loc(#loc5265)
    %18020 = stablehlo.sqrt %18019 : tensor<1x2048xf32> loc(#loc5266)
    %18021 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18022 = stablehlo.broadcast_in_dim %18021, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18023 = stablehlo.add %18020, %18022 : tensor<1x2048xf32> loc(#loc5265)
    %18024 = stablehlo.divide %15484, %18023 : tensor<1x2048xf32> loc(#loc5267)
    %18025 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18026 = stablehlo.broadcast_in_dim %18025, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %18027 = stablehlo.add %15562, %18026 : tensor<1x8192x2048xf32> loc(#loc5265)
    %18028 = stablehlo.sqrt %18027 : tensor<1x8192x2048xf32> loc(#loc5266)
    %18029 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18030 = stablehlo.broadcast_in_dim %18029, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %18031 = stablehlo.add %18028, %18030 : tensor<1x8192x2048xf32> loc(#loc5265)
    %18032 = stablehlo.divide %15551, %18031 : tensor<1x8192x2048xf32> loc(#loc5267)
    %18033 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18034 = stablehlo.broadcast_in_dim %18033, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18035 = stablehlo.add %15629, %18034 : tensor<1x2048xf32> loc(#loc5265)
    %18036 = stablehlo.sqrt %18035 : tensor<1x2048xf32> loc(#loc5266)
    %18037 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18038 = stablehlo.broadcast_in_dim %18037, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18039 = stablehlo.add %18036, %18038 : tensor<1x2048xf32> loc(#loc5265)
    %18040 = stablehlo.divide %15618, %18039 : tensor<1x2048xf32> loc(#loc5267)
    %18041 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18042 = stablehlo.broadcast_in_dim %18041, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18043 = stablehlo.add %15696, %18042 : tensor<1x2048xf32> loc(#loc5265)
    %18044 = stablehlo.sqrt %18043 : tensor<1x2048xf32> loc(#loc5266)
    %18045 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18046 = stablehlo.broadcast_in_dim %18045, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18047 = stablehlo.add %18044, %18046 : tensor<1x2048xf32> loc(#loc5265)
    %18048 = stablehlo.divide %15685, %18047 : tensor<1x2048xf32> loc(#loc5267)
    %18049 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18050 = stablehlo.broadcast_in_dim %18049, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18051 = stablehlo.add %15763, %18050 : tensor<1x2048xf32> loc(#loc5265)
    %18052 = stablehlo.sqrt %18051 : tensor<1x2048xf32> loc(#loc5266)
    %18053 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18054 = stablehlo.broadcast_in_dim %18053, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18055 = stablehlo.add %18052, %18054 : tensor<1x2048xf32> loc(#loc5265)
    %18056 = stablehlo.divide %15752, %18055 : tensor<1x2048xf32> loc(#loc5267)
    %18057 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18058 = stablehlo.broadcast_in_dim %18057, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18059 = stablehlo.add %15830, %18058 : tensor<1x2048xf32> loc(#loc5265)
    %18060 = stablehlo.sqrt %18059 : tensor<1x2048xf32> loc(#loc5266)
    %18061 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18062 = stablehlo.broadcast_in_dim %18061, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18063 = stablehlo.add %18060, %18062 : tensor<1x2048xf32> loc(#loc5265)
    %18064 = stablehlo.divide %15819, %18063 : tensor<1x2048xf32> loc(#loc5267)
    %18065 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18066 = stablehlo.broadcast_in_dim %18065, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18067 = stablehlo.add %15897, %18066 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18068 = stablehlo.sqrt %18067 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %18069 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18070 = stablehlo.broadcast_in_dim %18069, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18071 = stablehlo.add %18068, %18070 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18072 = stablehlo.divide %15886, %18071 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %18073 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18074 = stablehlo.broadcast_in_dim %18073, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %18075 = stablehlo.add %15964, %18074 : tensor<1x64xf32> loc(#loc5265)
    %18076 = stablehlo.sqrt %18075 : tensor<1x64xf32> loc(#loc5266)
    %18077 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18078 = stablehlo.broadcast_in_dim %18077, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %18079 = stablehlo.add %18076, %18078 : tensor<1x64xf32> loc(#loc5265)
    %18080 = stablehlo.divide %15953, %18079 : tensor<1x64xf32> loc(#loc5267)
    %18081 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18082 = stablehlo.broadcast_in_dim %18081, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %18083 = stablehlo.add %16031, %18082 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %18084 = stablehlo.sqrt %18083 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %18085 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18086 = stablehlo.broadcast_in_dim %18085, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %18087 = stablehlo.add %18084, %18086 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %18088 = stablehlo.divide %16020, %18087 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %18089 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18090 = stablehlo.broadcast_in_dim %18089, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %18091 = stablehlo.add %16098, %18090 : tensor<1x8192xf32> loc(#loc5265)
    %18092 = stablehlo.sqrt %18091 : tensor<1x8192xf32> loc(#loc5266)
    %18093 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18094 = stablehlo.broadcast_in_dim %18093, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5265)
    %18095 = stablehlo.add %18092, %18094 : tensor<1x8192xf32> loc(#loc5265)
    %18096 = stablehlo.divide %16087, %18095 : tensor<1x8192xf32> loc(#loc5267)
    %18097 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18098 = stablehlo.broadcast_in_dim %18097, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %18099 = stablehlo.add %16165, %18098 : tensor<1x2048x8192xf32> loc(#loc5265)
    %18100 = stablehlo.sqrt %18099 : tensor<1x2048x8192xf32> loc(#loc5266)
    %18101 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18102 = stablehlo.broadcast_in_dim %18101, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5265)
    %18103 = stablehlo.add %18100, %18102 : tensor<1x2048x8192xf32> loc(#loc5265)
    %18104 = stablehlo.divide %16154, %18103 : tensor<1x2048x8192xf32> loc(#loc5267)
    %18105 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18106 = stablehlo.broadcast_in_dim %18105, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18107 = stablehlo.add %16232, %18106 : tensor<1x2048xf32> loc(#loc5265)
    %18108 = stablehlo.sqrt %18107 : tensor<1x2048xf32> loc(#loc5266)
    %18109 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18110 = stablehlo.broadcast_in_dim %18109, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18111 = stablehlo.add %18108, %18110 : tensor<1x2048xf32> loc(#loc5265)
    %18112 = stablehlo.divide %16221, %18111 : tensor<1x2048xf32> loc(#loc5267)
    %18113 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18114 = stablehlo.broadcast_in_dim %18113, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %18115 = stablehlo.add %16299, %18114 : tensor<1x8192x2048xf32> loc(#loc5265)
    %18116 = stablehlo.sqrt %18115 : tensor<1x8192x2048xf32> loc(#loc5266)
    %18117 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18118 = stablehlo.broadcast_in_dim %18117, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5265)
    %18119 = stablehlo.add %18116, %18118 : tensor<1x8192x2048xf32> loc(#loc5265)
    %18120 = stablehlo.divide %16288, %18119 : tensor<1x8192x2048xf32> loc(#loc5267)
    %18121 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18122 = stablehlo.broadcast_in_dim %18121, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18123 = stablehlo.add %16366, %18122 : tensor<1x2048xf32> loc(#loc5265)
    %18124 = stablehlo.sqrt %18123 : tensor<1x2048xf32> loc(#loc5266)
    %18125 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18126 = stablehlo.broadcast_in_dim %18125, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18127 = stablehlo.add %18124, %18126 : tensor<1x2048xf32> loc(#loc5265)
    %18128 = stablehlo.divide %16355, %18127 : tensor<1x2048xf32> loc(#loc5267)
    %18129 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18130 = stablehlo.broadcast_in_dim %18129, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18131 = stablehlo.add %16433, %18130 : tensor<1x2048xf32> loc(#loc5265)
    %18132 = stablehlo.sqrt %18131 : tensor<1x2048xf32> loc(#loc5266)
    %18133 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18134 = stablehlo.broadcast_in_dim %18133, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18135 = stablehlo.add %18132, %18134 : tensor<1x2048xf32> loc(#loc5265)
    %18136 = stablehlo.divide %16422, %18135 : tensor<1x2048xf32> loc(#loc5267)
    %18137 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18138 = stablehlo.broadcast_in_dim %18137, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18139 = stablehlo.add %16500, %18138 : tensor<1x2048xf32> loc(#loc5265)
    %18140 = stablehlo.sqrt %18139 : tensor<1x2048xf32> loc(#loc5266)
    %18141 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18142 = stablehlo.broadcast_in_dim %18141, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18143 = stablehlo.add %18140, %18142 : tensor<1x2048xf32> loc(#loc5265)
    %18144 = stablehlo.divide %16489, %18143 : tensor<1x2048xf32> loc(#loc5267)
    %18145 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18146 = stablehlo.broadcast_in_dim %18145, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18147 = stablehlo.add %16567, %18146 : tensor<1x2048xf32> loc(#loc5265)
    %18148 = stablehlo.sqrt %18147 : tensor<1x2048xf32> loc(#loc5266)
    %18149 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18150 = stablehlo.broadcast_in_dim %18149, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5265)
    %18151 = stablehlo.add %18148, %18150 : tensor<1x2048xf32> loc(#loc5265)
    %18152 = stablehlo.divide %16556, %18151 : tensor<1x2048xf32> loc(#loc5267)
    %18153 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18154 = stablehlo.broadcast_in_dim %18153, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18155 = stablehlo.add %16634, %18154 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18156 = stablehlo.sqrt %18155 : tensor<1x3x2048x32x64xf32> loc(#loc5266)
    %18157 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18158 = stablehlo.broadcast_in_dim %18157, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18159 = stablehlo.add %18156, %18158 : tensor<1x3x2048x32x64xf32> loc(#loc5265)
    %18160 = stablehlo.divide %16623, %18159 : tensor<1x3x2048x32x64xf32> loc(#loc5267)
    %18161 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18162 = stablehlo.broadcast_in_dim %18161, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %18163 = stablehlo.add %16701, %18162 : tensor<1x64xf32> loc(#loc5265)
    %18164 = stablehlo.sqrt %18163 : tensor<1x64xf32> loc(#loc5266)
    %18165 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18166 = stablehlo.broadcast_in_dim %18165, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5265)
    %18167 = stablehlo.add %18164, %18166 : tensor<1x64xf32> loc(#loc5265)
    %18168 = stablehlo.divide %16690, %18167 : tensor<1x64xf32> loc(#loc5267)
    %18169 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %18170 = stablehlo.broadcast_in_dim %18169, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %18171 = stablehlo.add %16768, %18170 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %18172 = stablehlo.sqrt %18171 : tensor<1x2048x32x64xf32> loc(#loc5266)
    %18173 = stablehlo.constant dense<9.99999993E-9> : tensor<f32> loc(#loc)
    %18174 = stablehlo.broadcast_in_dim %18173, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5265)
    %18175 = stablehlo.add %18172, %18174 : tensor<1x2048x32x64xf32> loc(#loc5265)
    %18176 = stablehlo.divide %16757, %18175 : tensor<1x2048x32x64xf32> loc(#loc5267)
    %18177 = stablehlo.multiply %16776, %16776 : tensor<1x8192xf32> loc(#loc5268)
    %18178 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18179 = stablehlo.reduce(%18177 init: %18178) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18180 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %18181 = stablehlo.broadcast_in_dim %18180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18182 = stablehlo.divide %18179, %18181 : tensor<1xf32> loc(#loc5270)
    %18183 = stablehlo.sqrt %18182 : tensor<1xf32> loc(#loc5249)
    %18184 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18185 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18186 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18187 = call @nan_to_num_2078(%18183, %18184, %18185, %18186) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18189 = stablehlo.broadcast_in_dim %18188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18190 = stablehlo.divide %18187, %18189 : tensor<1xf32> loc(#loc5253)
    %18191 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18192 = stablehlo.broadcast_in_dim %18191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18193 = stablehlo.maximum %18192, %18190 : tensor<1xf32> loc(#loc5271)
    %18194 = stablehlo.broadcast_in_dim %18193, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18195 = stablehlo.broadcast_in_dim %18194, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %18196 = stablehlo.divide %16776, %18195 : tensor<1x8192xf32> loc(#loc5272)
    %18197 = stablehlo.multiply %16784, %16784 : tensor<1x2048x8192xf32> loc(#loc5268)
    %18198 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18199 = stablehlo.reduce(%18197 init: %18198) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18200 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18201 = stablehlo.broadcast_in_dim %18200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18202 = stablehlo.divide %18199, %18201 : tensor<1xf32> loc(#loc5270)
    %18203 = stablehlo.sqrt %18202 : tensor<1xf32> loc(#loc5249)
    %18204 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18205 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18206 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18207 = call @nan_to_num_2083(%18203, %18204, %18205, %18206) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18209 = stablehlo.broadcast_in_dim %18208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18210 = stablehlo.divide %18207, %18209 : tensor<1xf32> loc(#loc5253)
    %18211 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18212 = stablehlo.broadcast_in_dim %18211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18213 = stablehlo.maximum %18212, %18210 : tensor<1xf32> loc(#loc5271)
    %18214 = stablehlo.broadcast_in_dim %18213, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18215 = stablehlo.broadcast_in_dim %18214, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %18216 = stablehlo.divide %16784, %18215 : tensor<1x2048x8192xf32> loc(#loc5272)
    %18217 = stablehlo.multiply %16792, %16792 : tensor<1x2048xf32> loc(#loc5268)
    %18218 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18219 = stablehlo.reduce(%18217 init: %18218) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18220 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18221 = stablehlo.broadcast_in_dim %18220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18222 = stablehlo.divide %18219, %18221 : tensor<1xf32> loc(#loc5270)
    %18223 = stablehlo.sqrt %18222 : tensor<1xf32> loc(#loc5249)
    %18224 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18225 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18226 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18227 = call @nan_to_num_2088(%18223, %18224, %18225, %18226) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18228 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18229 = stablehlo.broadcast_in_dim %18228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18230 = stablehlo.divide %18227, %18229 : tensor<1xf32> loc(#loc5253)
    %18231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18232 = stablehlo.broadcast_in_dim %18231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18233 = stablehlo.maximum %18232, %18230 : tensor<1xf32> loc(#loc5271)
    %18234 = stablehlo.broadcast_in_dim %18233, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18235 = stablehlo.broadcast_in_dim %18234, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18236 = stablehlo.divide %16792, %18235 : tensor<1x2048xf32> loc(#loc5272)
    %18237 = stablehlo.multiply %16800, %16800 : tensor<1x8192x2048xf32> loc(#loc5268)
    %18238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18239 = stablehlo.reduce(%18237 init: %18238) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18240 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18241 = stablehlo.broadcast_in_dim %18240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18242 = stablehlo.divide %18239, %18241 : tensor<1xf32> loc(#loc5270)
    %18243 = stablehlo.sqrt %18242 : tensor<1xf32> loc(#loc5249)
    %18244 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18245 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18246 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18247 = call @nan_to_num_2093(%18243, %18244, %18245, %18246) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18248 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18249 = stablehlo.broadcast_in_dim %18248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18250 = stablehlo.divide %18247, %18249 : tensor<1xf32> loc(#loc5253)
    %18251 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18252 = stablehlo.broadcast_in_dim %18251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18253 = stablehlo.maximum %18252, %18250 : tensor<1xf32> loc(#loc5271)
    %18254 = stablehlo.broadcast_in_dim %18253, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18255 = stablehlo.broadcast_in_dim %18254, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %18256 = stablehlo.divide %16800, %18255 : tensor<1x8192x2048xf32> loc(#loc5272)
    %18257 = stablehlo.multiply %16808, %16808 : tensor<1x2048xf32> loc(#loc5268)
    %18258 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18259 = stablehlo.reduce(%18257 init: %18258) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18260 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18261 = stablehlo.broadcast_in_dim %18260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18262 = stablehlo.divide %18259, %18261 : tensor<1xf32> loc(#loc5270)
    %18263 = stablehlo.sqrt %18262 : tensor<1xf32> loc(#loc5249)
    %18264 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18265 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18266 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18267 = call @nan_to_num_2098(%18263, %18264, %18265, %18266) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18268 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18269 = stablehlo.broadcast_in_dim %18268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18270 = stablehlo.divide %18267, %18269 : tensor<1xf32> loc(#loc5253)
    %18271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18272 = stablehlo.broadcast_in_dim %18271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18273 = stablehlo.maximum %18272, %18270 : tensor<1xf32> loc(#loc5271)
    %18274 = stablehlo.broadcast_in_dim %18273, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18275 = stablehlo.broadcast_in_dim %18274, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18276 = stablehlo.divide %16808, %18275 : tensor<1x2048xf32> loc(#loc5272)
    %18277 = stablehlo.multiply %16816, %16816 : tensor<1x2048xf32> loc(#loc5268)
    %18278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18279 = stablehlo.reduce(%18277 init: %18278) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18280 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18281 = stablehlo.broadcast_in_dim %18280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18282 = stablehlo.divide %18279, %18281 : tensor<1xf32> loc(#loc5270)
    %18283 = stablehlo.sqrt %18282 : tensor<1xf32> loc(#loc5249)
    %18284 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18285 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18286 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18287 = call @nan_to_num_2103(%18283, %18284, %18285, %18286) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18288 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18289 = stablehlo.broadcast_in_dim %18288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18290 = stablehlo.divide %18287, %18289 : tensor<1xf32> loc(#loc5253)
    %18291 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18292 = stablehlo.broadcast_in_dim %18291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18293 = stablehlo.maximum %18292, %18290 : tensor<1xf32> loc(#loc5271)
    %18294 = stablehlo.broadcast_in_dim %18293, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18295 = stablehlo.broadcast_in_dim %18294, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18296 = stablehlo.divide %16816, %18295 : tensor<1x2048xf32> loc(#loc5272)
    %18297 = stablehlo.multiply %16824, %16824 : tensor<1x2048xf32> loc(#loc5268)
    %18298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18299 = stablehlo.reduce(%18297 init: %18298) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18300 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18301 = stablehlo.broadcast_in_dim %18300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18302 = stablehlo.divide %18299, %18301 : tensor<1xf32> loc(#loc5270)
    %18303 = stablehlo.sqrt %18302 : tensor<1xf32> loc(#loc5249)
    %18304 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18305 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18306 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18307 = call @nan_to_num_2108(%18303, %18304, %18305, %18306) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18308 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18309 = stablehlo.broadcast_in_dim %18308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18310 = stablehlo.divide %18307, %18309 : tensor<1xf32> loc(#loc5253)
    %18311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18312 = stablehlo.broadcast_in_dim %18311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18313 = stablehlo.maximum %18312, %18310 : tensor<1xf32> loc(#loc5271)
    %18314 = stablehlo.broadcast_in_dim %18313, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18315 = stablehlo.broadcast_in_dim %18314, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18316 = stablehlo.divide %16824, %18315 : tensor<1x2048xf32> loc(#loc5272)
    %18317 = stablehlo.multiply %16832, %16832 : tensor<1x2048xf32> loc(#loc5268)
    %18318 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18319 = stablehlo.reduce(%18317 init: %18318) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18320 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18321 = stablehlo.broadcast_in_dim %18320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18322 = stablehlo.divide %18319, %18321 : tensor<1xf32> loc(#loc5270)
    %18323 = stablehlo.sqrt %18322 : tensor<1xf32> loc(#loc5249)
    %18324 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18325 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18326 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18327 = call @nan_to_num_2113(%18323, %18324, %18325, %18326) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18328 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18329 = stablehlo.broadcast_in_dim %18328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18330 = stablehlo.divide %18327, %18329 : tensor<1xf32> loc(#loc5253)
    %18331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18332 = stablehlo.broadcast_in_dim %18331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18333 = stablehlo.maximum %18332, %18330 : tensor<1xf32> loc(#loc5271)
    %18334 = stablehlo.broadcast_in_dim %18333, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18335 = stablehlo.broadcast_in_dim %18334, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18336 = stablehlo.divide %16832, %18335 : tensor<1x2048xf32> loc(#loc5272)
    %18337 = stablehlo.multiply %16840, %16840 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %18338 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %18339 = stablehlo.reduce(%18337 init: %18338) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %18340 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %18341 = stablehlo.broadcast_in_dim %18340, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18342 = stablehlo.divide %18339, %18341 : tensor<1xf32> loc(#loc5270)
    %18343 = stablehlo.sqrt %18342 : tensor<1xf32> loc(#loc5249)
    %18344 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18345 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18346 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18347 = call @nan_to_num_2118(%18343, %18344, %18345, %18346) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18348 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18349 = stablehlo.broadcast_in_dim %18348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18350 = stablehlo.divide %18347, %18349 : tensor<1xf32> loc(#loc5253)
    %18351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18352 = stablehlo.broadcast_in_dim %18351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18353 = stablehlo.maximum %18352, %18350 : tensor<1xf32> loc(#loc5271)
    %18354 = stablehlo.broadcast_in_dim %18353, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %18355 = stablehlo.broadcast_in_dim %18354, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %18356 = stablehlo.divide %16840, %18355 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %18357 = stablehlo.multiply %16848, %16848 : tensor<1x64xf32> loc(#loc5268)
    %18358 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18359 = stablehlo.reduce(%18357 init: %18358) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18360 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %18361 = stablehlo.broadcast_in_dim %18360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18362 = stablehlo.divide %18359, %18361 : tensor<1xf32> loc(#loc5270)
    %18363 = stablehlo.sqrt %18362 : tensor<1xf32> loc(#loc5249)
    %18364 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18365 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18366 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18367 = call @nan_to_num_2123(%18363, %18364, %18365, %18366) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18369 = stablehlo.broadcast_in_dim %18368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18370 = stablehlo.divide %18367, %18369 : tensor<1xf32> loc(#loc5253)
    %18371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18372 = stablehlo.broadcast_in_dim %18371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18373 = stablehlo.maximum %18372, %18370 : tensor<1xf32> loc(#loc5271)
    %18374 = stablehlo.broadcast_in_dim %18373, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18375 = stablehlo.broadcast_in_dim %18374, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %18376 = stablehlo.divide %16848, %18375 : tensor<1x64xf32> loc(#loc5272)
    %18377 = stablehlo.multiply %16856, %16856 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %18378 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %18379 = stablehlo.reduce(%18377 init: %18378) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %18380 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %18381 = stablehlo.broadcast_in_dim %18380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18382 = stablehlo.divide %18379, %18381 : tensor<1xf32> loc(#loc5270)
    %18383 = stablehlo.sqrt %18382 : tensor<1xf32> loc(#loc5249)
    %18384 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18385 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18386 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18387 = call @nan_to_num_2128(%18383, %18384, %18385, %18386) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18389 = stablehlo.broadcast_in_dim %18388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18390 = stablehlo.divide %18387, %18389 : tensor<1xf32> loc(#loc5253)
    %18391 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18392 = stablehlo.broadcast_in_dim %18391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18393 = stablehlo.maximum %18392, %18390 : tensor<1xf32> loc(#loc5271)
    %18394 = stablehlo.broadcast_in_dim %18393, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %18395 = stablehlo.broadcast_in_dim %18394, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %18396 = stablehlo.divide %16856, %18395 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %18397 = stablehlo.multiply %16864, %16864 : tensor<1x8192xf32> loc(#loc5268)
    %18398 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18399 = stablehlo.reduce(%18397 init: %18398) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18400 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %18401 = stablehlo.broadcast_in_dim %18400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18402 = stablehlo.divide %18399, %18401 : tensor<1xf32> loc(#loc5270)
    %18403 = stablehlo.sqrt %18402 : tensor<1xf32> loc(#loc5249)
    %18404 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18405 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18406 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18407 = call @nan_to_num_2133(%18403, %18404, %18405, %18406) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18409 = stablehlo.broadcast_in_dim %18408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18410 = stablehlo.divide %18407, %18409 : tensor<1xf32> loc(#loc5253)
    %18411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18412 = stablehlo.broadcast_in_dim %18411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18413 = stablehlo.maximum %18412, %18410 : tensor<1xf32> loc(#loc5271)
    %18414 = stablehlo.broadcast_in_dim %18413, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18415 = stablehlo.broadcast_in_dim %18414, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %18416 = stablehlo.divide %16864, %18415 : tensor<1x8192xf32> loc(#loc5272)
    %18417 = stablehlo.multiply %16872, %16872 : tensor<1x2048x8192xf32> loc(#loc5268)
    %18418 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18419 = stablehlo.reduce(%18417 init: %18418) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18420 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18421 = stablehlo.broadcast_in_dim %18420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18422 = stablehlo.divide %18419, %18421 : tensor<1xf32> loc(#loc5270)
    %18423 = stablehlo.sqrt %18422 : tensor<1xf32> loc(#loc5249)
    %18424 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18425 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18426 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18427 = call @nan_to_num_2138(%18423, %18424, %18425, %18426) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18429 = stablehlo.broadcast_in_dim %18428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18430 = stablehlo.divide %18427, %18429 : tensor<1xf32> loc(#loc5253)
    %18431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18432 = stablehlo.broadcast_in_dim %18431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18433 = stablehlo.maximum %18432, %18430 : tensor<1xf32> loc(#loc5271)
    %18434 = stablehlo.broadcast_in_dim %18433, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18435 = stablehlo.broadcast_in_dim %18434, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %18436 = stablehlo.divide %16872, %18435 : tensor<1x2048x8192xf32> loc(#loc5272)
    %18437 = stablehlo.multiply %16880, %16880 : tensor<1x2048xf32> loc(#loc5268)
    %18438 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18439 = stablehlo.reduce(%18437 init: %18438) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18440 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18441 = stablehlo.broadcast_in_dim %18440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18442 = stablehlo.divide %18439, %18441 : tensor<1xf32> loc(#loc5270)
    %18443 = stablehlo.sqrt %18442 : tensor<1xf32> loc(#loc5249)
    %18444 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18445 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18446 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18447 = call @nan_to_num_2143(%18443, %18444, %18445, %18446) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18448 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18449 = stablehlo.broadcast_in_dim %18448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18450 = stablehlo.divide %18447, %18449 : tensor<1xf32> loc(#loc5253)
    %18451 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18452 = stablehlo.broadcast_in_dim %18451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18453 = stablehlo.maximum %18452, %18450 : tensor<1xf32> loc(#loc5271)
    %18454 = stablehlo.broadcast_in_dim %18453, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18455 = stablehlo.broadcast_in_dim %18454, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18456 = stablehlo.divide %16880, %18455 : tensor<1x2048xf32> loc(#loc5272)
    %18457 = stablehlo.multiply %16888, %16888 : tensor<1x8192x2048xf32> loc(#loc5268)
    %18458 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18459 = stablehlo.reduce(%18457 init: %18458) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18460 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18461 = stablehlo.broadcast_in_dim %18460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18462 = stablehlo.divide %18459, %18461 : tensor<1xf32> loc(#loc5270)
    %18463 = stablehlo.sqrt %18462 : tensor<1xf32> loc(#loc5249)
    %18464 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18465 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18466 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18467 = call @nan_to_num_2148(%18463, %18464, %18465, %18466) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18468 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18469 = stablehlo.broadcast_in_dim %18468, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18470 = stablehlo.divide %18467, %18469 : tensor<1xf32> loc(#loc5253)
    %18471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18472 = stablehlo.broadcast_in_dim %18471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18473 = stablehlo.maximum %18472, %18470 : tensor<1xf32> loc(#loc5271)
    %18474 = stablehlo.broadcast_in_dim %18473, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18475 = stablehlo.broadcast_in_dim %18474, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %18476 = stablehlo.divide %16888, %18475 : tensor<1x8192x2048xf32> loc(#loc5272)
    %18477 = stablehlo.multiply %16896, %16896 : tensor<1x2048xf32> loc(#loc5268)
    %18478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18479 = stablehlo.reduce(%18477 init: %18478) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18480 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18481 = stablehlo.broadcast_in_dim %18480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18482 = stablehlo.divide %18479, %18481 : tensor<1xf32> loc(#loc5270)
    %18483 = stablehlo.sqrt %18482 : tensor<1xf32> loc(#loc5249)
    %18484 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18485 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18486 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18487 = call @nan_to_num_2153(%18483, %18484, %18485, %18486) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18488 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18489 = stablehlo.broadcast_in_dim %18488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18490 = stablehlo.divide %18487, %18489 : tensor<1xf32> loc(#loc5253)
    %18491 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18492 = stablehlo.broadcast_in_dim %18491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18493 = stablehlo.maximum %18492, %18490 : tensor<1xf32> loc(#loc5271)
    %18494 = stablehlo.broadcast_in_dim %18493, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18495 = stablehlo.broadcast_in_dim %18494, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18496 = stablehlo.divide %16896, %18495 : tensor<1x2048xf32> loc(#loc5272)
    %18497 = stablehlo.multiply %16904, %16904 : tensor<1x2048xf32> loc(#loc5268)
    %18498 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18499 = stablehlo.reduce(%18497 init: %18498) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18500 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18501 = stablehlo.broadcast_in_dim %18500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18502 = stablehlo.divide %18499, %18501 : tensor<1xf32> loc(#loc5270)
    %18503 = stablehlo.sqrt %18502 : tensor<1xf32> loc(#loc5249)
    %18504 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18505 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18506 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18507 = call @nan_to_num_2158(%18503, %18504, %18505, %18506) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18508 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18509 = stablehlo.broadcast_in_dim %18508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18510 = stablehlo.divide %18507, %18509 : tensor<1xf32> loc(#loc5253)
    %18511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18512 = stablehlo.broadcast_in_dim %18511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18513 = stablehlo.maximum %18512, %18510 : tensor<1xf32> loc(#loc5271)
    %18514 = stablehlo.broadcast_in_dim %18513, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18515 = stablehlo.broadcast_in_dim %18514, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18516 = stablehlo.divide %16904, %18515 : tensor<1x2048xf32> loc(#loc5272)
    %18517 = stablehlo.multiply %16912, %16912 : tensor<1x2048xf32> loc(#loc5268)
    %18518 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18519 = stablehlo.reduce(%18517 init: %18518) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18520 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18521 = stablehlo.broadcast_in_dim %18520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18522 = stablehlo.divide %18519, %18521 : tensor<1xf32> loc(#loc5270)
    %18523 = stablehlo.sqrt %18522 : tensor<1xf32> loc(#loc5249)
    %18524 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18525 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18526 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18527 = call @nan_to_num_2163(%18523, %18524, %18525, %18526) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18529 = stablehlo.broadcast_in_dim %18528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18530 = stablehlo.divide %18527, %18529 : tensor<1xf32> loc(#loc5253)
    %18531 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18532 = stablehlo.broadcast_in_dim %18531, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18533 = stablehlo.maximum %18532, %18530 : tensor<1xf32> loc(#loc5271)
    %18534 = stablehlo.broadcast_in_dim %18533, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18535 = stablehlo.broadcast_in_dim %18534, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18536 = stablehlo.divide %16912, %18535 : tensor<1x2048xf32> loc(#loc5272)
    %18537 = stablehlo.multiply %16920, %16920 : tensor<1x2048xf32> loc(#loc5268)
    %18538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18539 = stablehlo.reduce(%18537 init: %18538) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18540 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18541 = stablehlo.broadcast_in_dim %18540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18542 = stablehlo.divide %18539, %18541 : tensor<1xf32> loc(#loc5270)
    %18543 = stablehlo.sqrt %18542 : tensor<1xf32> loc(#loc5249)
    %18544 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18545 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18546 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18547 = call @nan_to_num_2168(%18543, %18544, %18545, %18546) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18548 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18549 = stablehlo.broadcast_in_dim %18548, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18550 = stablehlo.divide %18547, %18549 : tensor<1xf32> loc(#loc5253)
    %18551 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18552 = stablehlo.broadcast_in_dim %18551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18553 = stablehlo.maximum %18552, %18550 : tensor<1xf32> loc(#loc5271)
    %18554 = stablehlo.broadcast_in_dim %18553, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18555 = stablehlo.broadcast_in_dim %18554, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18556 = stablehlo.divide %16920, %18555 : tensor<1x2048xf32> loc(#loc5272)
    %18557 = stablehlo.multiply %16928, %16928 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %18558 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %18559 = stablehlo.reduce(%18557 init: %18558) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %18560 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %18561 = stablehlo.broadcast_in_dim %18560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18562 = stablehlo.divide %18559, %18561 : tensor<1xf32> loc(#loc5270)
    %18563 = stablehlo.sqrt %18562 : tensor<1xf32> loc(#loc5249)
    %18564 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18565 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18566 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18567 = call @nan_to_num_2173(%18563, %18564, %18565, %18566) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18568 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18569 = stablehlo.broadcast_in_dim %18568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18570 = stablehlo.divide %18567, %18569 : tensor<1xf32> loc(#loc5253)
    %18571 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18572 = stablehlo.broadcast_in_dim %18571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18573 = stablehlo.maximum %18572, %18570 : tensor<1xf32> loc(#loc5271)
    %18574 = stablehlo.broadcast_in_dim %18573, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %18575 = stablehlo.broadcast_in_dim %18574, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %18576 = stablehlo.divide %16928, %18575 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %18577 = stablehlo.multiply %16936, %16936 : tensor<1x64xf32> loc(#loc5268)
    %18578 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18579 = stablehlo.reduce(%18577 init: %18578) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18580 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %18581 = stablehlo.broadcast_in_dim %18580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18582 = stablehlo.divide %18579, %18581 : tensor<1xf32> loc(#loc5270)
    %18583 = stablehlo.sqrt %18582 : tensor<1xf32> loc(#loc5249)
    %18584 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18585 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18586 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18587 = call @nan_to_num_2178(%18583, %18584, %18585, %18586) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18588 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18589 = stablehlo.broadcast_in_dim %18588, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18590 = stablehlo.divide %18587, %18589 : tensor<1xf32> loc(#loc5253)
    %18591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18592 = stablehlo.broadcast_in_dim %18591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18593 = stablehlo.maximum %18592, %18590 : tensor<1xf32> loc(#loc5271)
    %18594 = stablehlo.broadcast_in_dim %18593, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18595 = stablehlo.broadcast_in_dim %18594, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %18596 = stablehlo.divide %16936, %18595 : tensor<1x64xf32> loc(#loc5272)
    %18597 = stablehlo.multiply %16944, %16944 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %18598 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %18599 = stablehlo.reduce(%18597 init: %18598) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %18600 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %18601 = stablehlo.broadcast_in_dim %18600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18602 = stablehlo.divide %18599, %18601 : tensor<1xf32> loc(#loc5270)
    %18603 = stablehlo.sqrt %18602 : tensor<1xf32> loc(#loc5249)
    %18604 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18605 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18606 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18607 = call @nan_to_num_2183(%18603, %18604, %18605, %18606) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18608 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18609 = stablehlo.broadcast_in_dim %18608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18610 = stablehlo.divide %18607, %18609 : tensor<1xf32> loc(#loc5253)
    %18611 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18612 = stablehlo.broadcast_in_dim %18611, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18613 = stablehlo.maximum %18612, %18610 : tensor<1xf32> loc(#loc5271)
    %18614 = stablehlo.broadcast_in_dim %18613, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %18615 = stablehlo.broadcast_in_dim %18614, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %18616 = stablehlo.divide %16944, %18615 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %18617 = stablehlo.multiply %16952, %16952 : tensor<1x8192xf32> loc(#loc5268)
    %18618 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18619 = stablehlo.reduce(%18617 init: %18618) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18620 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %18621 = stablehlo.broadcast_in_dim %18620, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18622 = stablehlo.divide %18619, %18621 : tensor<1xf32> loc(#loc5270)
    %18623 = stablehlo.sqrt %18622 : tensor<1xf32> loc(#loc5249)
    %18624 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18625 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18626 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18627 = call @nan_to_num_2188(%18623, %18624, %18625, %18626) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18628 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18629 = stablehlo.broadcast_in_dim %18628, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18630 = stablehlo.divide %18627, %18629 : tensor<1xf32> loc(#loc5253)
    %18631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18632 = stablehlo.broadcast_in_dim %18631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18633 = stablehlo.maximum %18632, %18630 : tensor<1xf32> loc(#loc5271)
    %18634 = stablehlo.broadcast_in_dim %18633, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18635 = stablehlo.broadcast_in_dim %18634, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %18636 = stablehlo.divide %16952, %18635 : tensor<1x8192xf32> loc(#loc5272)
    %18637 = stablehlo.multiply %16960, %16960 : tensor<1x2048x8192xf32> loc(#loc5268)
    %18638 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18639 = stablehlo.reduce(%18637 init: %18638) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18640 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18641 = stablehlo.broadcast_in_dim %18640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18642 = stablehlo.divide %18639, %18641 : tensor<1xf32> loc(#loc5270)
    %18643 = stablehlo.sqrt %18642 : tensor<1xf32> loc(#loc5249)
    %18644 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18645 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18646 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18647 = call @nan_to_num_2193(%18643, %18644, %18645, %18646) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18648 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18649 = stablehlo.broadcast_in_dim %18648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18650 = stablehlo.divide %18647, %18649 : tensor<1xf32> loc(#loc5253)
    %18651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18652 = stablehlo.broadcast_in_dim %18651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18653 = stablehlo.maximum %18652, %18650 : tensor<1xf32> loc(#loc5271)
    %18654 = stablehlo.broadcast_in_dim %18653, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18655 = stablehlo.broadcast_in_dim %18654, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %18656 = stablehlo.divide %16960, %18655 : tensor<1x2048x8192xf32> loc(#loc5272)
    %18657 = stablehlo.multiply %16968, %16968 : tensor<1x2048xf32> loc(#loc5268)
    %18658 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18659 = stablehlo.reduce(%18657 init: %18658) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18660 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18661 = stablehlo.broadcast_in_dim %18660, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18662 = stablehlo.divide %18659, %18661 : tensor<1xf32> loc(#loc5270)
    %18663 = stablehlo.sqrt %18662 : tensor<1xf32> loc(#loc5249)
    %18664 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18665 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18666 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18667 = call @nan_to_num_2198(%18663, %18664, %18665, %18666) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18668 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18669 = stablehlo.broadcast_in_dim %18668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18670 = stablehlo.divide %18667, %18669 : tensor<1xf32> loc(#loc5253)
    %18671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18672 = stablehlo.broadcast_in_dim %18671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18673 = stablehlo.maximum %18672, %18670 : tensor<1xf32> loc(#loc5271)
    %18674 = stablehlo.broadcast_in_dim %18673, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18675 = stablehlo.broadcast_in_dim %18674, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18676 = stablehlo.divide %16968, %18675 : tensor<1x2048xf32> loc(#loc5272)
    %18677 = stablehlo.multiply %16976, %16976 : tensor<1x8192x2048xf32> loc(#loc5268)
    %18678 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18679 = stablehlo.reduce(%18677 init: %18678) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18680 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18681 = stablehlo.broadcast_in_dim %18680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18682 = stablehlo.divide %18679, %18681 : tensor<1xf32> loc(#loc5270)
    %18683 = stablehlo.sqrt %18682 : tensor<1xf32> loc(#loc5249)
    %18684 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18685 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18686 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18687 = call @nan_to_num_2203(%18683, %18684, %18685, %18686) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18688 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18689 = stablehlo.broadcast_in_dim %18688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18690 = stablehlo.divide %18687, %18689 : tensor<1xf32> loc(#loc5253)
    %18691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18692 = stablehlo.broadcast_in_dim %18691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18693 = stablehlo.maximum %18692, %18690 : tensor<1xf32> loc(#loc5271)
    %18694 = stablehlo.broadcast_in_dim %18693, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18695 = stablehlo.broadcast_in_dim %18694, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %18696 = stablehlo.divide %16976, %18695 : tensor<1x8192x2048xf32> loc(#loc5272)
    %18697 = stablehlo.multiply %16984, %16984 : tensor<1x2048xf32> loc(#loc5268)
    %18698 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18699 = stablehlo.reduce(%18697 init: %18698) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18700 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18701 = stablehlo.broadcast_in_dim %18700, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18702 = stablehlo.divide %18699, %18701 : tensor<1xf32> loc(#loc5270)
    %18703 = stablehlo.sqrt %18702 : tensor<1xf32> loc(#loc5249)
    %18704 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18705 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18706 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18707 = call @nan_to_num_2208(%18703, %18704, %18705, %18706) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18708 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18709 = stablehlo.broadcast_in_dim %18708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18710 = stablehlo.divide %18707, %18709 : tensor<1xf32> loc(#loc5253)
    %18711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18712 = stablehlo.broadcast_in_dim %18711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18713 = stablehlo.maximum %18712, %18710 : tensor<1xf32> loc(#loc5271)
    %18714 = stablehlo.broadcast_in_dim %18713, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18715 = stablehlo.broadcast_in_dim %18714, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18716 = stablehlo.divide %16984, %18715 : tensor<1x2048xf32> loc(#loc5272)
    %18717 = stablehlo.multiply %16992, %16992 : tensor<1x2048xf32> loc(#loc5268)
    %18718 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18719 = stablehlo.reduce(%18717 init: %18718) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18720 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18721 = stablehlo.broadcast_in_dim %18720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18722 = stablehlo.divide %18719, %18721 : tensor<1xf32> loc(#loc5270)
    %18723 = stablehlo.sqrt %18722 : tensor<1xf32> loc(#loc5249)
    %18724 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18725 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18726 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18727 = call @nan_to_num_2213(%18723, %18724, %18725, %18726) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18728 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18729 = stablehlo.broadcast_in_dim %18728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18730 = stablehlo.divide %18727, %18729 : tensor<1xf32> loc(#loc5253)
    %18731 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18732 = stablehlo.broadcast_in_dim %18731, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18733 = stablehlo.maximum %18732, %18730 : tensor<1xf32> loc(#loc5271)
    %18734 = stablehlo.broadcast_in_dim %18733, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18735 = stablehlo.broadcast_in_dim %18734, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18736 = stablehlo.divide %16992, %18735 : tensor<1x2048xf32> loc(#loc5272)
    %18737 = stablehlo.multiply %17000, %17000 : tensor<1x2048xf32> loc(#loc5268)
    %18738 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18739 = stablehlo.reduce(%18737 init: %18738) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18740 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18741 = stablehlo.broadcast_in_dim %18740, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18742 = stablehlo.divide %18739, %18741 : tensor<1xf32> loc(#loc5270)
    %18743 = stablehlo.sqrt %18742 : tensor<1xf32> loc(#loc5249)
    %18744 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18745 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18746 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18747 = call @nan_to_num_2218(%18743, %18744, %18745, %18746) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18749 = stablehlo.broadcast_in_dim %18748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18750 = stablehlo.divide %18747, %18749 : tensor<1xf32> loc(#loc5253)
    %18751 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18752 = stablehlo.broadcast_in_dim %18751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18753 = stablehlo.maximum %18752, %18750 : tensor<1xf32> loc(#loc5271)
    %18754 = stablehlo.broadcast_in_dim %18753, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18755 = stablehlo.broadcast_in_dim %18754, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18756 = stablehlo.divide %17000, %18755 : tensor<1x2048xf32> loc(#loc5272)
    %18757 = stablehlo.multiply %17008, %17008 : tensor<1x2048xf32> loc(#loc5268)
    %18758 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18759 = stablehlo.reduce(%18757 init: %18758) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18760 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18761 = stablehlo.broadcast_in_dim %18760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18762 = stablehlo.divide %18759, %18761 : tensor<1xf32> loc(#loc5270)
    %18763 = stablehlo.sqrt %18762 : tensor<1xf32> loc(#loc5249)
    %18764 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18765 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18766 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18767 = call @nan_to_num_2223(%18763, %18764, %18765, %18766) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18768 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18769 = stablehlo.broadcast_in_dim %18768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18770 = stablehlo.divide %18767, %18769 : tensor<1xf32> loc(#loc5253)
    %18771 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18772 = stablehlo.broadcast_in_dim %18771, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18773 = stablehlo.maximum %18772, %18770 : tensor<1xf32> loc(#loc5271)
    %18774 = stablehlo.broadcast_in_dim %18773, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18775 = stablehlo.broadcast_in_dim %18774, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18776 = stablehlo.divide %17008, %18775 : tensor<1x2048xf32> loc(#loc5272)
    %18777 = stablehlo.multiply %17016, %17016 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %18778 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %18779 = stablehlo.reduce(%18777 init: %18778) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %18780 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %18781 = stablehlo.broadcast_in_dim %18780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18782 = stablehlo.divide %18779, %18781 : tensor<1xf32> loc(#loc5270)
    %18783 = stablehlo.sqrt %18782 : tensor<1xf32> loc(#loc5249)
    %18784 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18785 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18786 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18787 = call @nan_to_num_2228(%18783, %18784, %18785, %18786) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18788 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18789 = stablehlo.broadcast_in_dim %18788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18790 = stablehlo.divide %18787, %18789 : tensor<1xf32> loc(#loc5253)
    %18791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18792 = stablehlo.broadcast_in_dim %18791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18793 = stablehlo.maximum %18792, %18790 : tensor<1xf32> loc(#loc5271)
    %18794 = stablehlo.broadcast_in_dim %18793, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %18795 = stablehlo.broadcast_in_dim %18794, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %18796 = stablehlo.divide %17016, %18795 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %18797 = stablehlo.multiply %17024, %17024 : tensor<1x64xf32> loc(#loc5268)
    %18798 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18799 = stablehlo.reduce(%18797 init: %18798) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18800 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %18801 = stablehlo.broadcast_in_dim %18800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18802 = stablehlo.divide %18799, %18801 : tensor<1xf32> loc(#loc5270)
    %18803 = stablehlo.sqrt %18802 : tensor<1xf32> loc(#loc5249)
    %18804 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18805 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18806 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18807 = call @nan_to_num_2233(%18803, %18804, %18805, %18806) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18808 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18809 = stablehlo.broadcast_in_dim %18808, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18810 = stablehlo.divide %18807, %18809 : tensor<1xf32> loc(#loc5253)
    %18811 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18812 = stablehlo.broadcast_in_dim %18811, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18813 = stablehlo.maximum %18812, %18810 : tensor<1xf32> loc(#loc5271)
    %18814 = stablehlo.broadcast_in_dim %18813, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18815 = stablehlo.broadcast_in_dim %18814, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %18816 = stablehlo.divide %17024, %18815 : tensor<1x64xf32> loc(#loc5272)
    %18817 = stablehlo.multiply %17032, %17032 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %18818 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %18819 = stablehlo.reduce(%18817 init: %18818) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %18820 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %18821 = stablehlo.broadcast_in_dim %18820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18822 = stablehlo.divide %18819, %18821 : tensor<1xf32> loc(#loc5270)
    %18823 = stablehlo.sqrt %18822 : tensor<1xf32> loc(#loc5249)
    %18824 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18825 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18826 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18827 = call @nan_to_num_2238(%18823, %18824, %18825, %18826) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18828 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18829 = stablehlo.broadcast_in_dim %18828, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18830 = stablehlo.divide %18827, %18829 : tensor<1xf32> loc(#loc5253)
    %18831 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18832 = stablehlo.broadcast_in_dim %18831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18833 = stablehlo.maximum %18832, %18830 : tensor<1xf32> loc(#loc5271)
    %18834 = stablehlo.broadcast_in_dim %18833, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %18835 = stablehlo.broadcast_in_dim %18834, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %18836 = stablehlo.divide %17032, %18835 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %18837 = stablehlo.multiply %17040, %17040 : tensor<1x8192xf32> loc(#loc5268)
    %18838 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18839 = stablehlo.reduce(%18837 init: %18838) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18840 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %18841 = stablehlo.broadcast_in_dim %18840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18842 = stablehlo.divide %18839, %18841 : tensor<1xf32> loc(#loc5270)
    %18843 = stablehlo.sqrt %18842 : tensor<1xf32> loc(#loc5249)
    %18844 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18845 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18846 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18847 = call @nan_to_num_2243(%18843, %18844, %18845, %18846) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18848 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18849 = stablehlo.broadcast_in_dim %18848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18850 = stablehlo.divide %18847, %18849 : tensor<1xf32> loc(#loc5253)
    %18851 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18852 = stablehlo.broadcast_in_dim %18851, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18853 = stablehlo.maximum %18852, %18850 : tensor<1xf32> loc(#loc5271)
    %18854 = stablehlo.broadcast_in_dim %18853, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18855 = stablehlo.broadcast_in_dim %18854, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %18856 = stablehlo.divide %17040, %18855 : tensor<1x8192xf32> loc(#loc5272)
    %18857 = stablehlo.multiply %17048, %17048 : tensor<1x2048x8192xf32> loc(#loc5268)
    %18858 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18859 = stablehlo.reduce(%18857 init: %18858) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18860 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18861 = stablehlo.broadcast_in_dim %18860, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18862 = stablehlo.divide %18859, %18861 : tensor<1xf32> loc(#loc5270)
    %18863 = stablehlo.sqrt %18862 : tensor<1xf32> loc(#loc5249)
    %18864 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18865 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18866 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18867 = call @nan_to_num_2248(%18863, %18864, %18865, %18866) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18868 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18869 = stablehlo.broadcast_in_dim %18868, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18870 = stablehlo.divide %18867, %18869 : tensor<1xf32> loc(#loc5253)
    %18871 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18872 = stablehlo.broadcast_in_dim %18871, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18873 = stablehlo.maximum %18872, %18870 : tensor<1xf32> loc(#loc5271)
    %18874 = stablehlo.broadcast_in_dim %18873, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18875 = stablehlo.broadcast_in_dim %18874, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %18876 = stablehlo.divide %17048, %18875 : tensor<1x2048x8192xf32> loc(#loc5272)
    %18877 = stablehlo.multiply %17056, %17056 : tensor<1x2048xf32> loc(#loc5268)
    %18878 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18879 = stablehlo.reduce(%18877 init: %18878) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18880 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18881 = stablehlo.broadcast_in_dim %18880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18882 = stablehlo.divide %18879, %18881 : tensor<1xf32> loc(#loc5270)
    %18883 = stablehlo.sqrt %18882 : tensor<1xf32> loc(#loc5249)
    %18884 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18885 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18886 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18887 = call @nan_to_num_2253(%18883, %18884, %18885, %18886) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18888 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18889 = stablehlo.broadcast_in_dim %18888, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18890 = stablehlo.divide %18887, %18889 : tensor<1xf32> loc(#loc5253)
    %18891 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18892 = stablehlo.broadcast_in_dim %18891, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18893 = stablehlo.maximum %18892, %18890 : tensor<1xf32> loc(#loc5271)
    %18894 = stablehlo.broadcast_in_dim %18893, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18895 = stablehlo.broadcast_in_dim %18894, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18896 = stablehlo.divide %17056, %18895 : tensor<1x2048xf32> loc(#loc5272)
    %18897 = stablehlo.multiply %17064, %17064 : tensor<1x8192x2048xf32> loc(#loc5268)
    %18898 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %18899 = stablehlo.reduce(%18897 init: %18898) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %18900 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %18901 = stablehlo.broadcast_in_dim %18900, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18902 = stablehlo.divide %18899, %18901 : tensor<1xf32> loc(#loc5270)
    %18903 = stablehlo.sqrt %18902 : tensor<1xf32> loc(#loc5249)
    %18904 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18905 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18906 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18907 = call @nan_to_num_2258(%18903, %18904, %18905, %18906) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18908 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18909 = stablehlo.broadcast_in_dim %18908, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18910 = stablehlo.divide %18907, %18909 : tensor<1xf32> loc(#loc5253)
    %18911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18912 = stablehlo.broadcast_in_dim %18911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18913 = stablehlo.maximum %18912, %18910 : tensor<1xf32> loc(#loc5271)
    %18914 = stablehlo.broadcast_in_dim %18913, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %18915 = stablehlo.broadcast_in_dim %18914, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %18916 = stablehlo.divide %17064, %18915 : tensor<1x8192x2048xf32> loc(#loc5272)
    %18917 = stablehlo.multiply %17072, %17072 : tensor<1x2048xf32> loc(#loc5268)
    %18918 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18919 = stablehlo.reduce(%18917 init: %18918) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18920 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18921 = stablehlo.broadcast_in_dim %18920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18922 = stablehlo.divide %18919, %18921 : tensor<1xf32> loc(#loc5270)
    %18923 = stablehlo.sqrt %18922 : tensor<1xf32> loc(#loc5249)
    %18924 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18925 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18926 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18927 = call @nan_to_num_2263(%18923, %18924, %18925, %18926) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18928 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18929 = stablehlo.broadcast_in_dim %18928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18930 = stablehlo.divide %18927, %18929 : tensor<1xf32> loc(#loc5253)
    %18931 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18932 = stablehlo.broadcast_in_dim %18931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18933 = stablehlo.maximum %18932, %18930 : tensor<1xf32> loc(#loc5271)
    %18934 = stablehlo.broadcast_in_dim %18933, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18935 = stablehlo.broadcast_in_dim %18934, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18936 = stablehlo.divide %17072, %18935 : tensor<1x2048xf32> loc(#loc5272)
    %18937 = stablehlo.multiply %17080, %17080 : tensor<1x2048xf32> loc(#loc5268)
    %18938 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18939 = stablehlo.reduce(%18937 init: %18938) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18940 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18941 = stablehlo.broadcast_in_dim %18940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18942 = stablehlo.divide %18939, %18941 : tensor<1xf32> loc(#loc5270)
    %18943 = stablehlo.sqrt %18942 : tensor<1xf32> loc(#loc5249)
    %18944 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18945 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18946 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18947 = call @nan_to_num_2268(%18943, %18944, %18945, %18946) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18948 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18949 = stablehlo.broadcast_in_dim %18948, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18950 = stablehlo.divide %18947, %18949 : tensor<1xf32> loc(#loc5253)
    %18951 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18952 = stablehlo.broadcast_in_dim %18951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18953 = stablehlo.maximum %18952, %18950 : tensor<1xf32> loc(#loc5271)
    %18954 = stablehlo.broadcast_in_dim %18953, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18955 = stablehlo.broadcast_in_dim %18954, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18956 = stablehlo.divide %17080, %18955 : tensor<1x2048xf32> loc(#loc5272)
    %18957 = stablehlo.multiply %17088, %17088 : tensor<1x2048xf32> loc(#loc5268)
    %18958 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18959 = stablehlo.reduce(%18957 init: %18958) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18960 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18961 = stablehlo.broadcast_in_dim %18960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18962 = stablehlo.divide %18959, %18961 : tensor<1xf32> loc(#loc5270)
    %18963 = stablehlo.sqrt %18962 : tensor<1xf32> loc(#loc5249)
    %18964 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18965 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18966 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18967 = call @nan_to_num_2273(%18963, %18964, %18965, %18966) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18969 = stablehlo.broadcast_in_dim %18968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18970 = stablehlo.divide %18967, %18969 : tensor<1xf32> loc(#loc5253)
    %18971 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18972 = stablehlo.broadcast_in_dim %18971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18973 = stablehlo.maximum %18972, %18970 : tensor<1xf32> loc(#loc5271)
    %18974 = stablehlo.broadcast_in_dim %18973, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18975 = stablehlo.broadcast_in_dim %18974, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18976 = stablehlo.divide %17088, %18975 : tensor<1x2048xf32> loc(#loc5272)
    %18977 = stablehlo.multiply %17096, %17096 : tensor<1x2048xf32> loc(#loc5268)
    %18978 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %18979 = stablehlo.reduce(%18977 init: %18978) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %18980 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %18981 = stablehlo.broadcast_in_dim %18980, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %18982 = stablehlo.divide %18979, %18981 : tensor<1xf32> loc(#loc5270)
    %18983 = stablehlo.sqrt %18982 : tensor<1xf32> loc(#loc5249)
    %18984 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18985 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18986 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %18987 = call @nan_to_num_2278(%18983, %18984, %18985, %18986) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %18988 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18989 = stablehlo.broadcast_in_dim %18988, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %18990 = stablehlo.divide %18987, %18989 : tensor<1xf32> loc(#loc5253)
    %18991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %18992 = stablehlo.broadcast_in_dim %18991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %18993 = stablehlo.maximum %18992, %18990 : tensor<1xf32> loc(#loc5271)
    %18994 = stablehlo.broadcast_in_dim %18993, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %18995 = stablehlo.broadcast_in_dim %18994, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %18996 = stablehlo.divide %17096, %18995 : tensor<1x2048xf32> loc(#loc5272)
    %18997 = stablehlo.multiply %17104, %17104 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %18998 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %18999 = stablehlo.reduce(%18997 init: %18998) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %19000 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %19001 = stablehlo.broadcast_in_dim %19000, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19002 = stablehlo.divide %18999, %19001 : tensor<1xf32> loc(#loc5270)
    %19003 = stablehlo.sqrt %19002 : tensor<1xf32> loc(#loc5249)
    %19004 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19005 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19006 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19007 = call @nan_to_num_2283(%19003, %19004, %19005, %19006) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19008 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19009 = stablehlo.broadcast_in_dim %19008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19010 = stablehlo.divide %19007, %19009 : tensor<1xf32> loc(#loc5253)
    %19011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19012 = stablehlo.broadcast_in_dim %19011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19013 = stablehlo.maximum %19012, %19010 : tensor<1xf32> loc(#loc5271)
    %19014 = stablehlo.broadcast_in_dim %19013, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %19015 = stablehlo.broadcast_in_dim %19014, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19016 = stablehlo.divide %17104, %19015 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19017 = stablehlo.multiply %17112, %17112 : tensor<1x64xf32> loc(#loc5268)
    %19018 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19019 = stablehlo.reduce(%19017 init: %19018) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19020 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %19021 = stablehlo.broadcast_in_dim %19020, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19022 = stablehlo.divide %19019, %19021 : tensor<1xf32> loc(#loc5270)
    %19023 = stablehlo.sqrt %19022 : tensor<1xf32> loc(#loc5249)
    %19024 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19025 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19026 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19027 = call @nan_to_num_2288(%19023, %19024, %19025, %19026) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19028 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19029 = stablehlo.broadcast_in_dim %19028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19030 = stablehlo.divide %19027, %19029 : tensor<1xf32> loc(#loc5253)
    %19031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19032 = stablehlo.broadcast_in_dim %19031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19033 = stablehlo.maximum %19032, %19030 : tensor<1xf32> loc(#loc5271)
    %19034 = stablehlo.broadcast_in_dim %19033, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19035 = stablehlo.broadcast_in_dim %19034, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %19036 = stablehlo.divide %17112, %19035 : tensor<1x64xf32> loc(#loc5272)
    %19037 = stablehlo.multiply %17120, %17120 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %19038 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %19039 = stablehlo.reduce(%19037 init: %19038) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %19040 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %19041 = stablehlo.broadcast_in_dim %19040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19042 = stablehlo.divide %19039, %19041 : tensor<1xf32> loc(#loc5270)
    %19043 = stablehlo.sqrt %19042 : tensor<1xf32> loc(#loc5249)
    %19044 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19045 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19046 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19047 = call @nan_to_num_2293(%19043, %19044, %19045, %19046) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19048 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19049 = stablehlo.broadcast_in_dim %19048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19050 = stablehlo.divide %19047, %19049 : tensor<1xf32> loc(#loc5253)
    %19051 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19052 = stablehlo.broadcast_in_dim %19051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19053 = stablehlo.maximum %19052, %19050 : tensor<1xf32> loc(#loc5271)
    %19054 = stablehlo.broadcast_in_dim %19053, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %19055 = stablehlo.broadcast_in_dim %19054, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %19056 = stablehlo.divide %17120, %19055 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %19057 = stablehlo.multiply %17128, %17128 : tensor<1x8192xf32> loc(#loc5268)
    %19058 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19059 = stablehlo.reduce(%19057 init: %19058) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19060 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %19061 = stablehlo.broadcast_in_dim %19060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19062 = stablehlo.divide %19059, %19061 : tensor<1xf32> loc(#loc5270)
    %19063 = stablehlo.sqrt %19062 : tensor<1xf32> loc(#loc5249)
    %19064 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19065 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19066 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19067 = call @nan_to_num_2298(%19063, %19064, %19065, %19066) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19069 = stablehlo.broadcast_in_dim %19068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19070 = stablehlo.divide %19067, %19069 : tensor<1xf32> loc(#loc5253)
    %19071 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19072 = stablehlo.broadcast_in_dim %19071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19073 = stablehlo.maximum %19072, %19070 : tensor<1xf32> loc(#loc5271)
    %19074 = stablehlo.broadcast_in_dim %19073, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19075 = stablehlo.broadcast_in_dim %19074, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %19076 = stablehlo.divide %17128, %19075 : tensor<1x8192xf32> loc(#loc5272)
    %19077 = stablehlo.multiply %17136, %17136 : tensor<1x2048x8192xf32> loc(#loc5268)
    %19078 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19079 = stablehlo.reduce(%19077 init: %19078) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19080 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19081 = stablehlo.broadcast_in_dim %19080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19082 = stablehlo.divide %19079, %19081 : tensor<1xf32> loc(#loc5270)
    %19083 = stablehlo.sqrt %19082 : tensor<1xf32> loc(#loc5249)
    %19084 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19085 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19086 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19087 = call @nan_to_num_2303(%19083, %19084, %19085, %19086) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19089 = stablehlo.broadcast_in_dim %19088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19090 = stablehlo.divide %19087, %19089 : tensor<1xf32> loc(#loc5253)
    %19091 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19092 = stablehlo.broadcast_in_dim %19091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19093 = stablehlo.maximum %19092, %19090 : tensor<1xf32> loc(#loc5271)
    %19094 = stablehlo.broadcast_in_dim %19093, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19095 = stablehlo.broadcast_in_dim %19094, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %19096 = stablehlo.divide %17136, %19095 : tensor<1x2048x8192xf32> loc(#loc5272)
    %19097 = stablehlo.multiply %17144, %17144 : tensor<1x2048xf32> loc(#loc5268)
    %19098 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19099 = stablehlo.reduce(%19097 init: %19098) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19100 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19101 = stablehlo.broadcast_in_dim %19100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19102 = stablehlo.divide %19099, %19101 : tensor<1xf32> loc(#loc5270)
    %19103 = stablehlo.sqrt %19102 : tensor<1xf32> loc(#loc5249)
    %19104 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19105 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19106 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19107 = call @nan_to_num_2308(%19103, %19104, %19105, %19106) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19108 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19109 = stablehlo.broadcast_in_dim %19108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19110 = stablehlo.divide %19107, %19109 : tensor<1xf32> loc(#loc5253)
    %19111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19112 = stablehlo.broadcast_in_dim %19111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19113 = stablehlo.maximum %19112, %19110 : tensor<1xf32> loc(#loc5271)
    %19114 = stablehlo.broadcast_in_dim %19113, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19115 = stablehlo.broadcast_in_dim %19114, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19116 = stablehlo.divide %17144, %19115 : tensor<1x2048xf32> loc(#loc5272)
    %19117 = stablehlo.multiply %17152, %17152 : tensor<1x8192x2048xf32> loc(#loc5268)
    %19118 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19119 = stablehlo.reduce(%19117 init: %19118) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19120 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19121 = stablehlo.broadcast_in_dim %19120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19122 = stablehlo.divide %19119, %19121 : tensor<1xf32> loc(#loc5270)
    %19123 = stablehlo.sqrt %19122 : tensor<1xf32> loc(#loc5249)
    %19124 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19125 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19126 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19127 = call @nan_to_num_2313(%19123, %19124, %19125, %19126) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19128 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19129 = stablehlo.broadcast_in_dim %19128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19130 = stablehlo.divide %19127, %19129 : tensor<1xf32> loc(#loc5253)
    %19131 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19132 = stablehlo.broadcast_in_dim %19131, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19133 = stablehlo.maximum %19132, %19130 : tensor<1xf32> loc(#loc5271)
    %19134 = stablehlo.broadcast_in_dim %19133, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19135 = stablehlo.broadcast_in_dim %19134, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %19136 = stablehlo.divide %17152, %19135 : tensor<1x8192x2048xf32> loc(#loc5272)
    %19137 = stablehlo.multiply %17160, %17160 : tensor<1x2048xf32> loc(#loc5268)
    %19138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19139 = stablehlo.reduce(%19137 init: %19138) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19140 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19141 = stablehlo.broadcast_in_dim %19140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19142 = stablehlo.divide %19139, %19141 : tensor<1xf32> loc(#loc5270)
    %19143 = stablehlo.sqrt %19142 : tensor<1xf32> loc(#loc5249)
    %19144 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19145 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19146 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19147 = call @nan_to_num_2318(%19143, %19144, %19145, %19146) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19148 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19149 = stablehlo.broadcast_in_dim %19148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19150 = stablehlo.divide %19147, %19149 : tensor<1xf32> loc(#loc5253)
    %19151 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19152 = stablehlo.broadcast_in_dim %19151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19153 = stablehlo.maximum %19152, %19150 : tensor<1xf32> loc(#loc5271)
    %19154 = stablehlo.broadcast_in_dim %19153, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19155 = stablehlo.broadcast_in_dim %19154, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19156 = stablehlo.divide %17160, %19155 : tensor<1x2048xf32> loc(#loc5272)
    %19157 = stablehlo.multiply %17168, %17168 : tensor<1x2048xf32> loc(#loc5268)
    %19158 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19159 = stablehlo.reduce(%19157 init: %19158) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19160 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19161 = stablehlo.broadcast_in_dim %19160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19162 = stablehlo.divide %19159, %19161 : tensor<1xf32> loc(#loc5270)
    %19163 = stablehlo.sqrt %19162 : tensor<1xf32> loc(#loc5249)
    %19164 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19165 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19166 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19167 = call @nan_to_num_2323(%19163, %19164, %19165, %19166) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19168 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19169 = stablehlo.broadcast_in_dim %19168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19170 = stablehlo.divide %19167, %19169 : tensor<1xf32> loc(#loc5253)
    %19171 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19172 = stablehlo.broadcast_in_dim %19171, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19173 = stablehlo.maximum %19172, %19170 : tensor<1xf32> loc(#loc5271)
    %19174 = stablehlo.broadcast_in_dim %19173, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19175 = stablehlo.broadcast_in_dim %19174, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19176 = stablehlo.divide %17168, %19175 : tensor<1x2048xf32> loc(#loc5272)
    %19177 = stablehlo.multiply %17176, %17176 : tensor<1x2048xf32> loc(#loc5268)
    %19178 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19179 = stablehlo.reduce(%19177 init: %19178) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19180 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19181 = stablehlo.broadcast_in_dim %19180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19182 = stablehlo.divide %19179, %19181 : tensor<1xf32> loc(#loc5270)
    %19183 = stablehlo.sqrt %19182 : tensor<1xf32> loc(#loc5249)
    %19184 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19185 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19186 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19187 = call @nan_to_num_2328(%19183, %19184, %19185, %19186) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19189 = stablehlo.broadcast_in_dim %19188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19190 = stablehlo.divide %19187, %19189 : tensor<1xf32> loc(#loc5253)
    %19191 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19192 = stablehlo.broadcast_in_dim %19191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19193 = stablehlo.maximum %19192, %19190 : tensor<1xf32> loc(#loc5271)
    %19194 = stablehlo.broadcast_in_dim %19193, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19195 = stablehlo.broadcast_in_dim %19194, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19196 = stablehlo.divide %17176, %19195 : tensor<1x2048xf32> loc(#loc5272)
    %19197 = stablehlo.multiply %17184, %17184 : tensor<1x2048xf32> loc(#loc5268)
    %19198 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19199 = stablehlo.reduce(%19197 init: %19198) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19200 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19201 = stablehlo.broadcast_in_dim %19200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19202 = stablehlo.divide %19199, %19201 : tensor<1xf32> loc(#loc5270)
    %19203 = stablehlo.sqrt %19202 : tensor<1xf32> loc(#loc5249)
    %19204 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19205 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19206 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19207 = call @nan_to_num_2333(%19203, %19204, %19205, %19206) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19209 = stablehlo.broadcast_in_dim %19208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19210 = stablehlo.divide %19207, %19209 : tensor<1xf32> loc(#loc5253)
    %19211 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19212 = stablehlo.broadcast_in_dim %19211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19213 = stablehlo.maximum %19212, %19210 : tensor<1xf32> loc(#loc5271)
    %19214 = stablehlo.broadcast_in_dim %19213, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19215 = stablehlo.broadcast_in_dim %19214, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19216 = stablehlo.divide %17184, %19215 : tensor<1x2048xf32> loc(#loc5272)
    %19217 = stablehlo.multiply %17192, %17192 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %19218 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %19219 = stablehlo.reduce(%19217 init: %19218) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %19220 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %19221 = stablehlo.broadcast_in_dim %19220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19222 = stablehlo.divide %19219, %19221 : tensor<1xf32> loc(#loc5270)
    %19223 = stablehlo.sqrt %19222 : tensor<1xf32> loc(#loc5249)
    %19224 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19225 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19226 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19227 = call @nan_to_num_2338(%19223, %19224, %19225, %19226) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19228 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19229 = stablehlo.broadcast_in_dim %19228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19230 = stablehlo.divide %19227, %19229 : tensor<1xf32> loc(#loc5253)
    %19231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19232 = stablehlo.broadcast_in_dim %19231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19233 = stablehlo.maximum %19232, %19230 : tensor<1xf32> loc(#loc5271)
    %19234 = stablehlo.broadcast_in_dim %19233, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %19235 = stablehlo.broadcast_in_dim %19234, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19236 = stablehlo.divide %17192, %19235 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19237 = stablehlo.multiply %17200, %17200 : tensor<1x64xf32> loc(#loc5268)
    %19238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19239 = stablehlo.reduce(%19237 init: %19238) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19240 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %19241 = stablehlo.broadcast_in_dim %19240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19242 = stablehlo.divide %19239, %19241 : tensor<1xf32> loc(#loc5270)
    %19243 = stablehlo.sqrt %19242 : tensor<1xf32> loc(#loc5249)
    %19244 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19245 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19246 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19247 = call @nan_to_num_2343(%19243, %19244, %19245, %19246) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19248 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19249 = stablehlo.broadcast_in_dim %19248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19250 = stablehlo.divide %19247, %19249 : tensor<1xf32> loc(#loc5253)
    %19251 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19252 = stablehlo.broadcast_in_dim %19251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19253 = stablehlo.maximum %19252, %19250 : tensor<1xf32> loc(#loc5271)
    %19254 = stablehlo.broadcast_in_dim %19253, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19255 = stablehlo.broadcast_in_dim %19254, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %19256 = stablehlo.divide %17200, %19255 : tensor<1x64xf32> loc(#loc5272)
    %19257 = stablehlo.multiply %17208, %17208 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %19258 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %19259 = stablehlo.reduce(%19257 init: %19258) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %19260 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %19261 = stablehlo.broadcast_in_dim %19260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19262 = stablehlo.divide %19259, %19261 : tensor<1xf32> loc(#loc5270)
    %19263 = stablehlo.sqrt %19262 : tensor<1xf32> loc(#loc5249)
    %19264 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19265 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19266 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19267 = call @nan_to_num_2348(%19263, %19264, %19265, %19266) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19268 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19269 = stablehlo.broadcast_in_dim %19268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19270 = stablehlo.divide %19267, %19269 : tensor<1xf32> loc(#loc5253)
    %19271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19272 = stablehlo.broadcast_in_dim %19271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19273 = stablehlo.maximum %19272, %19270 : tensor<1xf32> loc(#loc5271)
    %19274 = stablehlo.broadcast_in_dim %19273, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %19275 = stablehlo.broadcast_in_dim %19274, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %19276 = stablehlo.divide %17208, %19275 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %19277 = stablehlo.multiply %17216, %17216 : tensor<1x8192xf32> loc(#loc5268)
    %19278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19279 = stablehlo.reduce(%19277 init: %19278) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19280 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %19281 = stablehlo.broadcast_in_dim %19280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19282 = stablehlo.divide %19279, %19281 : tensor<1xf32> loc(#loc5270)
    %19283 = stablehlo.sqrt %19282 : tensor<1xf32> loc(#loc5249)
    %19284 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19285 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19286 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19287 = call @nan_to_num_2353(%19283, %19284, %19285, %19286) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19288 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19289 = stablehlo.broadcast_in_dim %19288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19290 = stablehlo.divide %19287, %19289 : tensor<1xf32> loc(#loc5253)
    %19291 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19292 = stablehlo.broadcast_in_dim %19291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19293 = stablehlo.maximum %19292, %19290 : tensor<1xf32> loc(#loc5271)
    %19294 = stablehlo.broadcast_in_dim %19293, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19295 = stablehlo.broadcast_in_dim %19294, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %19296 = stablehlo.divide %17216, %19295 : tensor<1x8192xf32> loc(#loc5272)
    %19297 = stablehlo.multiply %17224, %17224 : tensor<1x2048x8192xf32> loc(#loc5268)
    %19298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19299 = stablehlo.reduce(%19297 init: %19298) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19300 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19301 = stablehlo.broadcast_in_dim %19300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19302 = stablehlo.divide %19299, %19301 : tensor<1xf32> loc(#loc5270)
    %19303 = stablehlo.sqrt %19302 : tensor<1xf32> loc(#loc5249)
    %19304 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19305 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19306 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19307 = call @nan_to_num_2358(%19303, %19304, %19305, %19306) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19308 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19309 = stablehlo.broadcast_in_dim %19308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19310 = stablehlo.divide %19307, %19309 : tensor<1xf32> loc(#loc5253)
    %19311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19312 = stablehlo.broadcast_in_dim %19311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19313 = stablehlo.maximum %19312, %19310 : tensor<1xf32> loc(#loc5271)
    %19314 = stablehlo.broadcast_in_dim %19313, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19315 = stablehlo.broadcast_in_dim %19314, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %19316 = stablehlo.divide %17224, %19315 : tensor<1x2048x8192xf32> loc(#loc5272)
    %19317 = stablehlo.multiply %17232, %17232 : tensor<1x2048xf32> loc(#loc5268)
    %19318 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19319 = stablehlo.reduce(%19317 init: %19318) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19320 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19321 = stablehlo.broadcast_in_dim %19320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19322 = stablehlo.divide %19319, %19321 : tensor<1xf32> loc(#loc5270)
    %19323 = stablehlo.sqrt %19322 : tensor<1xf32> loc(#loc5249)
    %19324 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19325 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19326 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19327 = call @nan_to_num_2363(%19323, %19324, %19325, %19326) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19328 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19329 = stablehlo.broadcast_in_dim %19328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19330 = stablehlo.divide %19327, %19329 : tensor<1xf32> loc(#loc5253)
    %19331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19332 = stablehlo.broadcast_in_dim %19331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19333 = stablehlo.maximum %19332, %19330 : tensor<1xf32> loc(#loc5271)
    %19334 = stablehlo.broadcast_in_dim %19333, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19335 = stablehlo.broadcast_in_dim %19334, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19336 = stablehlo.divide %17232, %19335 : tensor<1x2048xf32> loc(#loc5272)
    %19337 = stablehlo.multiply %17240, %17240 : tensor<1x8192x2048xf32> loc(#loc5268)
    %19338 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19339 = stablehlo.reduce(%19337 init: %19338) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19340 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19341 = stablehlo.broadcast_in_dim %19340, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19342 = stablehlo.divide %19339, %19341 : tensor<1xf32> loc(#loc5270)
    %19343 = stablehlo.sqrt %19342 : tensor<1xf32> loc(#loc5249)
    %19344 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19345 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19346 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19347 = call @nan_to_num_2368(%19343, %19344, %19345, %19346) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19348 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19349 = stablehlo.broadcast_in_dim %19348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19350 = stablehlo.divide %19347, %19349 : tensor<1xf32> loc(#loc5253)
    %19351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19352 = stablehlo.broadcast_in_dim %19351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19353 = stablehlo.maximum %19352, %19350 : tensor<1xf32> loc(#loc5271)
    %19354 = stablehlo.broadcast_in_dim %19353, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19355 = stablehlo.broadcast_in_dim %19354, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %19356 = stablehlo.divide %17240, %19355 : tensor<1x8192x2048xf32> loc(#loc5272)
    %19357 = stablehlo.multiply %17248, %17248 : tensor<1x2048xf32> loc(#loc5268)
    %19358 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19359 = stablehlo.reduce(%19357 init: %19358) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19360 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19361 = stablehlo.broadcast_in_dim %19360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19362 = stablehlo.divide %19359, %19361 : tensor<1xf32> loc(#loc5270)
    %19363 = stablehlo.sqrt %19362 : tensor<1xf32> loc(#loc5249)
    %19364 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19365 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19366 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19367 = call @nan_to_num_2373(%19363, %19364, %19365, %19366) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19369 = stablehlo.broadcast_in_dim %19368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19370 = stablehlo.divide %19367, %19369 : tensor<1xf32> loc(#loc5253)
    %19371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19372 = stablehlo.broadcast_in_dim %19371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19373 = stablehlo.maximum %19372, %19370 : tensor<1xf32> loc(#loc5271)
    %19374 = stablehlo.broadcast_in_dim %19373, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19375 = stablehlo.broadcast_in_dim %19374, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19376 = stablehlo.divide %17248, %19375 : tensor<1x2048xf32> loc(#loc5272)
    %19377 = stablehlo.multiply %17256, %17256 : tensor<1x2048xf32> loc(#loc5268)
    %19378 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19379 = stablehlo.reduce(%19377 init: %19378) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19380 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19381 = stablehlo.broadcast_in_dim %19380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19382 = stablehlo.divide %19379, %19381 : tensor<1xf32> loc(#loc5270)
    %19383 = stablehlo.sqrt %19382 : tensor<1xf32> loc(#loc5249)
    %19384 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19385 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19386 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19387 = call @nan_to_num_2378(%19383, %19384, %19385, %19386) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19389 = stablehlo.broadcast_in_dim %19388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19390 = stablehlo.divide %19387, %19389 : tensor<1xf32> loc(#loc5253)
    %19391 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19392 = stablehlo.broadcast_in_dim %19391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19393 = stablehlo.maximum %19392, %19390 : tensor<1xf32> loc(#loc5271)
    %19394 = stablehlo.broadcast_in_dim %19393, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19395 = stablehlo.broadcast_in_dim %19394, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19396 = stablehlo.divide %17256, %19395 : tensor<1x2048xf32> loc(#loc5272)
    %19397 = stablehlo.multiply %17264, %17264 : tensor<1x2048xf32> loc(#loc5268)
    %19398 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19399 = stablehlo.reduce(%19397 init: %19398) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19400 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19401 = stablehlo.broadcast_in_dim %19400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19402 = stablehlo.divide %19399, %19401 : tensor<1xf32> loc(#loc5270)
    %19403 = stablehlo.sqrt %19402 : tensor<1xf32> loc(#loc5249)
    %19404 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19405 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19406 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19407 = call @nan_to_num_2383(%19403, %19404, %19405, %19406) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19409 = stablehlo.broadcast_in_dim %19408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19410 = stablehlo.divide %19407, %19409 : tensor<1xf32> loc(#loc5253)
    %19411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19412 = stablehlo.broadcast_in_dim %19411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19413 = stablehlo.maximum %19412, %19410 : tensor<1xf32> loc(#loc5271)
    %19414 = stablehlo.broadcast_in_dim %19413, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19415 = stablehlo.broadcast_in_dim %19414, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19416 = stablehlo.divide %17264, %19415 : tensor<1x2048xf32> loc(#loc5272)
    %19417 = stablehlo.multiply %17272, %17272 : tensor<1x2048xf32> loc(#loc5268)
    %19418 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19419 = stablehlo.reduce(%19417 init: %19418) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19420 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19421 = stablehlo.broadcast_in_dim %19420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19422 = stablehlo.divide %19419, %19421 : tensor<1xf32> loc(#loc5270)
    %19423 = stablehlo.sqrt %19422 : tensor<1xf32> loc(#loc5249)
    %19424 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19425 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19426 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19427 = call @nan_to_num_2388(%19423, %19424, %19425, %19426) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19429 = stablehlo.broadcast_in_dim %19428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19430 = stablehlo.divide %19427, %19429 : tensor<1xf32> loc(#loc5253)
    %19431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19432 = stablehlo.broadcast_in_dim %19431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19433 = stablehlo.maximum %19432, %19430 : tensor<1xf32> loc(#loc5271)
    %19434 = stablehlo.broadcast_in_dim %19433, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19435 = stablehlo.broadcast_in_dim %19434, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19436 = stablehlo.divide %17272, %19435 : tensor<1x2048xf32> loc(#loc5272)
    %19437 = stablehlo.multiply %17280, %17280 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %19438 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %19439 = stablehlo.reduce(%19437 init: %19438) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %19440 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %19441 = stablehlo.broadcast_in_dim %19440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19442 = stablehlo.divide %19439, %19441 : tensor<1xf32> loc(#loc5270)
    %19443 = stablehlo.sqrt %19442 : tensor<1xf32> loc(#loc5249)
    %19444 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19445 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19446 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19447 = call @nan_to_num_2393(%19443, %19444, %19445, %19446) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19448 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19449 = stablehlo.broadcast_in_dim %19448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19450 = stablehlo.divide %19447, %19449 : tensor<1xf32> loc(#loc5253)
    %19451 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19452 = stablehlo.broadcast_in_dim %19451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19453 = stablehlo.maximum %19452, %19450 : tensor<1xf32> loc(#loc5271)
    %19454 = stablehlo.broadcast_in_dim %19453, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %19455 = stablehlo.broadcast_in_dim %19454, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19456 = stablehlo.divide %17280, %19455 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19457 = stablehlo.multiply %17288, %17288 : tensor<1x64xf32> loc(#loc5268)
    %19458 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19459 = stablehlo.reduce(%19457 init: %19458) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19460 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %19461 = stablehlo.broadcast_in_dim %19460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19462 = stablehlo.divide %19459, %19461 : tensor<1xf32> loc(#loc5270)
    %19463 = stablehlo.sqrt %19462 : tensor<1xf32> loc(#loc5249)
    %19464 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19465 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19466 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19467 = call @nan_to_num_2398(%19463, %19464, %19465, %19466) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19468 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19469 = stablehlo.broadcast_in_dim %19468, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19470 = stablehlo.divide %19467, %19469 : tensor<1xf32> loc(#loc5253)
    %19471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19472 = stablehlo.broadcast_in_dim %19471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19473 = stablehlo.maximum %19472, %19470 : tensor<1xf32> loc(#loc5271)
    %19474 = stablehlo.broadcast_in_dim %19473, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19475 = stablehlo.broadcast_in_dim %19474, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %19476 = stablehlo.divide %17288, %19475 : tensor<1x64xf32> loc(#loc5272)
    %19477 = stablehlo.multiply %17296, %17296 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %19478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %19479 = stablehlo.reduce(%19477 init: %19478) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %19480 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %19481 = stablehlo.broadcast_in_dim %19480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19482 = stablehlo.divide %19479, %19481 : tensor<1xf32> loc(#loc5270)
    %19483 = stablehlo.sqrt %19482 : tensor<1xf32> loc(#loc5249)
    %19484 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19485 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19486 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19487 = call @nan_to_num_2403(%19483, %19484, %19485, %19486) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19488 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19489 = stablehlo.broadcast_in_dim %19488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19490 = stablehlo.divide %19487, %19489 : tensor<1xf32> loc(#loc5253)
    %19491 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19492 = stablehlo.broadcast_in_dim %19491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19493 = stablehlo.maximum %19492, %19490 : tensor<1xf32> loc(#loc5271)
    %19494 = stablehlo.broadcast_in_dim %19493, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %19495 = stablehlo.broadcast_in_dim %19494, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %19496 = stablehlo.divide %17296, %19495 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %19497 = stablehlo.multiply %17304, %17304 : tensor<1x8192xf32> loc(#loc5268)
    %19498 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19499 = stablehlo.reduce(%19497 init: %19498) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19500 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %19501 = stablehlo.broadcast_in_dim %19500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19502 = stablehlo.divide %19499, %19501 : tensor<1xf32> loc(#loc5270)
    %19503 = stablehlo.sqrt %19502 : tensor<1xf32> loc(#loc5249)
    %19504 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19505 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19506 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19507 = call @nan_to_num_2408(%19503, %19504, %19505, %19506) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19508 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19509 = stablehlo.broadcast_in_dim %19508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19510 = stablehlo.divide %19507, %19509 : tensor<1xf32> loc(#loc5253)
    %19511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19512 = stablehlo.broadcast_in_dim %19511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19513 = stablehlo.maximum %19512, %19510 : tensor<1xf32> loc(#loc5271)
    %19514 = stablehlo.broadcast_in_dim %19513, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19515 = stablehlo.broadcast_in_dim %19514, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %19516 = stablehlo.divide %17304, %19515 : tensor<1x8192xf32> loc(#loc5272)
    %19517 = stablehlo.multiply %17312, %17312 : tensor<1x2048x8192xf32> loc(#loc5268)
    %19518 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19519 = stablehlo.reduce(%19517 init: %19518) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19520 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19521 = stablehlo.broadcast_in_dim %19520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19522 = stablehlo.divide %19519, %19521 : tensor<1xf32> loc(#loc5270)
    %19523 = stablehlo.sqrt %19522 : tensor<1xf32> loc(#loc5249)
    %19524 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19525 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19526 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19527 = call @nan_to_num_2413(%19523, %19524, %19525, %19526) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19529 = stablehlo.broadcast_in_dim %19528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19530 = stablehlo.divide %19527, %19529 : tensor<1xf32> loc(#loc5253)
    %19531 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19532 = stablehlo.broadcast_in_dim %19531, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19533 = stablehlo.maximum %19532, %19530 : tensor<1xf32> loc(#loc5271)
    %19534 = stablehlo.broadcast_in_dim %19533, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19535 = stablehlo.broadcast_in_dim %19534, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %19536 = stablehlo.divide %17312, %19535 : tensor<1x2048x8192xf32> loc(#loc5272)
    %19537 = stablehlo.multiply %17320, %17320 : tensor<1x2048xf32> loc(#loc5268)
    %19538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19539 = stablehlo.reduce(%19537 init: %19538) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19540 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19541 = stablehlo.broadcast_in_dim %19540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19542 = stablehlo.divide %19539, %19541 : tensor<1xf32> loc(#loc5270)
    %19543 = stablehlo.sqrt %19542 : tensor<1xf32> loc(#loc5249)
    %19544 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19545 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19546 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19547 = call @nan_to_num_2418(%19543, %19544, %19545, %19546) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19548 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19549 = stablehlo.broadcast_in_dim %19548, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19550 = stablehlo.divide %19547, %19549 : tensor<1xf32> loc(#loc5253)
    %19551 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19552 = stablehlo.broadcast_in_dim %19551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19553 = stablehlo.maximum %19552, %19550 : tensor<1xf32> loc(#loc5271)
    %19554 = stablehlo.broadcast_in_dim %19553, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19555 = stablehlo.broadcast_in_dim %19554, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19556 = stablehlo.divide %17320, %19555 : tensor<1x2048xf32> loc(#loc5272)
    %19557 = stablehlo.multiply %17328, %17328 : tensor<1x8192x2048xf32> loc(#loc5268)
    %19558 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19559 = stablehlo.reduce(%19557 init: %19558) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19560 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19561 = stablehlo.broadcast_in_dim %19560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19562 = stablehlo.divide %19559, %19561 : tensor<1xf32> loc(#loc5270)
    %19563 = stablehlo.sqrt %19562 : tensor<1xf32> loc(#loc5249)
    %19564 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19565 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19566 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19567 = call @nan_to_num_2423(%19563, %19564, %19565, %19566) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19568 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19569 = stablehlo.broadcast_in_dim %19568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19570 = stablehlo.divide %19567, %19569 : tensor<1xf32> loc(#loc5253)
    %19571 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19572 = stablehlo.broadcast_in_dim %19571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19573 = stablehlo.maximum %19572, %19570 : tensor<1xf32> loc(#loc5271)
    %19574 = stablehlo.broadcast_in_dim %19573, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19575 = stablehlo.broadcast_in_dim %19574, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %19576 = stablehlo.divide %17328, %19575 : tensor<1x8192x2048xf32> loc(#loc5272)
    %19577 = stablehlo.multiply %17336, %17336 : tensor<1x2048xf32> loc(#loc5268)
    %19578 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19579 = stablehlo.reduce(%19577 init: %19578) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19580 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19581 = stablehlo.broadcast_in_dim %19580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19582 = stablehlo.divide %19579, %19581 : tensor<1xf32> loc(#loc5270)
    %19583 = stablehlo.sqrt %19582 : tensor<1xf32> loc(#loc5249)
    %19584 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19585 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19586 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19587 = call @nan_to_num_2428(%19583, %19584, %19585, %19586) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19588 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19589 = stablehlo.broadcast_in_dim %19588, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19590 = stablehlo.divide %19587, %19589 : tensor<1xf32> loc(#loc5253)
    %19591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19592 = stablehlo.broadcast_in_dim %19591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19593 = stablehlo.maximum %19592, %19590 : tensor<1xf32> loc(#loc5271)
    %19594 = stablehlo.broadcast_in_dim %19593, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19595 = stablehlo.broadcast_in_dim %19594, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19596 = stablehlo.divide %17336, %19595 : tensor<1x2048xf32> loc(#loc5272)
    %19597 = stablehlo.multiply %17344, %17344 : tensor<1x2048xf32> loc(#loc5268)
    %19598 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19599 = stablehlo.reduce(%19597 init: %19598) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19600 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19601 = stablehlo.broadcast_in_dim %19600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19602 = stablehlo.divide %19599, %19601 : tensor<1xf32> loc(#loc5270)
    %19603 = stablehlo.sqrt %19602 : tensor<1xf32> loc(#loc5249)
    %19604 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19605 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19606 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19607 = call @nan_to_num_2433(%19603, %19604, %19605, %19606) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19608 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19609 = stablehlo.broadcast_in_dim %19608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19610 = stablehlo.divide %19607, %19609 : tensor<1xf32> loc(#loc5253)
    %19611 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19612 = stablehlo.broadcast_in_dim %19611, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19613 = stablehlo.maximum %19612, %19610 : tensor<1xf32> loc(#loc5271)
    %19614 = stablehlo.broadcast_in_dim %19613, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19615 = stablehlo.broadcast_in_dim %19614, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19616 = stablehlo.divide %17344, %19615 : tensor<1x2048xf32> loc(#loc5272)
    %19617 = stablehlo.multiply %17352, %17352 : tensor<1x2048xf32> loc(#loc5268)
    %19618 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19619 = stablehlo.reduce(%19617 init: %19618) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19620 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19621 = stablehlo.broadcast_in_dim %19620, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19622 = stablehlo.divide %19619, %19621 : tensor<1xf32> loc(#loc5270)
    %19623 = stablehlo.sqrt %19622 : tensor<1xf32> loc(#loc5249)
    %19624 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19625 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19626 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19627 = call @nan_to_num_2438(%19623, %19624, %19625, %19626) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19628 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19629 = stablehlo.broadcast_in_dim %19628, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19630 = stablehlo.divide %19627, %19629 : tensor<1xf32> loc(#loc5253)
    %19631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19632 = stablehlo.broadcast_in_dim %19631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19633 = stablehlo.maximum %19632, %19630 : tensor<1xf32> loc(#loc5271)
    %19634 = stablehlo.broadcast_in_dim %19633, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19635 = stablehlo.broadcast_in_dim %19634, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19636 = stablehlo.divide %17352, %19635 : tensor<1x2048xf32> loc(#loc5272)
    %19637 = stablehlo.multiply %17360, %17360 : tensor<1x2048xf32> loc(#loc5268)
    %19638 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19639 = stablehlo.reduce(%19637 init: %19638) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19640 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19641 = stablehlo.broadcast_in_dim %19640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19642 = stablehlo.divide %19639, %19641 : tensor<1xf32> loc(#loc5270)
    %19643 = stablehlo.sqrt %19642 : tensor<1xf32> loc(#loc5249)
    %19644 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19645 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19646 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19647 = call @nan_to_num_2443(%19643, %19644, %19645, %19646) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19648 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19649 = stablehlo.broadcast_in_dim %19648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19650 = stablehlo.divide %19647, %19649 : tensor<1xf32> loc(#loc5253)
    %19651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19652 = stablehlo.broadcast_in_dim %19651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19653 = stablehlo.maximum %19652, %19650 : tensor<1xf32> loc(#loc5271)
    %19654 = stablehlo.broadcast_in_dim %19653, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19655 = stablehlo.broadcast_in_dim %19654, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19656 = stablehlo.divide %17360, %19655 : tensor<1x2048xf32> loc(#loc5272)
    %19657 = stablehlo.multiply %17368, %17368 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %19658 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %19659 = stablehlo.reduce(%19657 init: %19658) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %19660 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %19661 = stablehlo.broadcast_in_dim %19660, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19662 = stablehlo.divide %19659, %19661 : tensor<1xf32> loc(#loc5270)
    %19663 = stablehlo.sqrt %19662 : tensor<1xf32> loc(#loc5249)
    %19664 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19665 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19666 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19667 = call @nan_to_num_2448(%19663, %19664, %19665, %19666) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19668 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19669 = stablehlo.broadcast_in_dim %19668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19670 = stablehlo.divide %19667, %19669 : tensor<1xf32> loc(#loc5253)
    %19671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19672 = stablehlo.broadcast_in_dim %19671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19673 = stablehlo.maximum %19672, %19670 : tensor<1xf32> loc(#loc5271)
    %19674 = stablehlo.broadcast_in_dim %19673, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %19675 = stablehlo.broadcast_in_dim %19674, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19676 = stablehlo.divide %17368, %19675 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19677 = stablehlo.multiply %17376, %17376 : tensor<1x64xf32> loc(#loc5268)
    %19678 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19679 = stablehlo.reduce(%19677 init: %19678) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19680 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %19681 = stablehlo.broadcast_in_dim %19680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19682 = stablehlo.divide %19679, %19681 : tensor<1xf32> loc(#loc5270)
    %19683 = stablehlo.sqrt %19682 : tensor<1xf32> loc(#loc5249)
    %19684 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19685 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19686 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19687 = call @nan_to_num_2453(%19683, %19684, %19685, %19686) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19688 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19689 = stablehlo.broadcast_in_dim %19688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19690 = stablehlo.divide %19687, %19689 : tensor<1xf32> loc(#loc5253)
    %19691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19692 = stablehlo.broadcast_in_dim %19691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19693 = stablehlo.maximum %19692, %19690 : tensor<1xf32> loc(#loc5271)
    %19694 = stablehlo.broadcast_in_dim %19693, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19695 = stablehlo.broadcast_in_dim %19694, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %19696 = stablehlo.divide %17376, %19695 : tensor<1x64xf32> loc(#loc5272)
    %19697 = stablehlo.multiply %17384, %17384 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %19698 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %19699 = stablehlo.reduce(%19697 init: %19698) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %19700 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %19701 = stablehlo.broadcast_in_dim %19700, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19702 = stablehlo.divide %19699, %19701 : tensor<1xf32> loc(#loc5270)
    %19703 = stablehlo.sqrt %19702 : tensor<1xf32> loc(#loc5249)
    %19704 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19705 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19706 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19707 = call @nan_to_num_2458(%19703, %19704, %19705, %19706) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19708 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19709 = stablehlo.broadcast_in_dim %19708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19710 = stablehlo.divide %19707, %19709 : tensor<1xf32> loc(#loc5253)
    %19711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19712 = stablehlo.broadcast_in_dim %19711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19713 = stablehlo.maximum %19712, %19710 : tensor<1xf32> loc(#loc5271)
    %19714 = stablehlo.broadcast_in_dim %19713, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %19715 = stablehlo.broadcast_in_dim %19714, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %19716 = stablehlo.divide %17384, %19715 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %19717 = stablehlo.multiply %17392, %17392 : tensor<1x8192xf32> loc(#loc5268)
    %19718 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19719 = stablehlo.reduce(%19717 init: %19718) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19720 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %19721 = stablehlo.broadcast_in_dim %19720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19722 = stablehlo.divide %19719, %19721 : tensor<1xf32> loc(#loc5270)
    %19723 = stablehlo.sqrt %19722 : tensor<1xf32> loc(#loc5249)
    %19724 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19725 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19726 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19727 = call @nan_to_num_2463(%19723, %19724, %19725, %19726) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19728 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19729 = stablehlo.broadcast_in_dim %19728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19730 = stablehlo.divide %19727, %19729 : tensor<1xf32> loc(#loc5253)
    %19731 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19732 = stablehlo.broadcast_in_dim %19731, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19733 = stablehlo.maximum %19732, %19730 : tensor<1xf32> loc(#loc5271)
    %19734 = stablehlo.broadcast_in_dim %19733, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19735 = stablehlo.broadcast_in_dim %19734, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %19736 = stablehlo.divide %17392, %19735 : tensor<1x8192xf32> loc(#loc5272)
    %19737 = stablehlo.multiply %17400, %17400 : tensor<1x2048x8192xf32> loc(#loc5268)
    %19738 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19739 = stablehlo.reduce(%19737 init: %19738) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19740 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19741 = stablehlo.broadcast_in_dim %19740, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19742 = stablehlo.divide %19739, %19741 : tensor<1xf32> loc(#loc5270)
    %19743 = stablehlo.sqrt %19742 : tensor<1xf32> loc(#loc5249)
    %19744 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19745 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19746 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19747 = call @nan_to_num_2468(%19743, %19744, %19745, %19746) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19749 = stablehlo.broadcast_in_dim %19748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19750 = stablehlo.divide %19747, %19749 : tensor<1xf32> loc(#loc5253)
    %19751 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19752 = stablehlo.broadcast_in_dim %19751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19753 = stablehlo.maximum %19752, %19750 : tensor<1xf32> loc(#loc5271)
    %19754 = stablehlo.broadcast_in_dim %19753, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19755 = stablehlo.broadcast_in_dim %19754, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %19756 = stablehlo.divide %17400, %19755 : tensor<1x2048x8192xf32> loc(#loc5272)
    %19757 = stablehlo.multiply %17408, %17408 : tensor<1x2048xf32> loc(#loc5268)
    %19758 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19759 = stablehlo.reduce(%19757 init: %19758) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19760 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19761 = stablehlo.broadcast_in_dim %19760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19762 = stablehlo.divide %19759, %19761 : tensor<1xf32> loc(#loc5270)
    %19763 = stablehlo.sqrt %19762 : tensor<1xf32> loc(#loc5249)
    %19764 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19765 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19766 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19767 = call @nan_to_num_2473(%19763, %19764, %19765, %19766) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19768 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19769 = stablehlo.broadcast_in_dim %19768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19770 = stablehlo.divide %19767, %19769 : tensor<1xf32> loc(#loc5253)
    %19771 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19772 = stablehlo.broadcast_in_dim %19771, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19773 = stablehlo.maximum %19772, %19770 : tensor<1xf32> loc(#loc5271)
    %19774 = stablehlo.broadcast_in_dim %19773, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19775 = stablehlo.broadcast_in_dim %19774, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19776 = stablehlo.divide %17408, %19775 : tensor<1x2048xf32> loc(#loc5272)
    %19777 = stablehlo.multiply %17416, %17416 : tensor<1x8192x2048xf32> loc(#loc5268)
    %19778 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19779 = stablehlo.reduce(%19777 init: %19778) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19780 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19781 = stablehlo.broadcast_in_dim %19780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19782 = stablehlo.divide %19779, %19781 : tensor<1xf32> loc(#loc5270)
    %19783 = stablehlo.sqrt %19782 : tensor<1xf32> loc(#loc5249)
    %19784 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19785 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19786 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19787 = call @nan_to_num_2478(%19783, %19784, %19785, %19786) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19788 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19789 = stablehlo.broadcast_in_dim %19788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19790 = stablehlo.divide %19787, %19789 : tensor<1xf32> loc(#loc5253)
    %19791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19792 = stablehlo.broadcast_in_dim %19791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19793 = stablehlo.maximum %19792, %19790 : tensor<1xf32> loc(#loc5271)
    %19794 = stablehlo.broadcast_in_dim %19793, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19795 = stablehlo.broadcast_in_dim %19794, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %19796 = stablehlo.divide %17416, %19795 : tensor<1x8192x2048xf32> loc(#loc5272)
    %19797 = stablehlo.multiply %17424, %17424 : tensor<1x2048xf32> loc(#loc5268)
    %19798 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19799 = stablehlo.reduce(%19797 init: %19798) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19800 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19801 = stablehlo.broadcast_in_dim %19800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19802 = stablehlo.divide %19799, %19801 : tensor<1xf32> loc(#loc5270)
    %19803 = stablehlo.sqrt %19802 : tensor<1xf32> loc(#loc5249)
    %19804 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19805 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19806 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19807 = call @nan_to_num_2483(%19803, %19804, %19805, %19806) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19808 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19809 = stablehlo.broadcast_in_dim %19808, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19810 = stablehlo.divide %19807, %19809 : tensor<1xf32> loc(#loc5253)
    %19811 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19812 = stablehlo.broadcast_in_dim %19811, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19813 = stablehlo.maximum %19812, %19810 : tensor<1xf32> loc(#loc5271)
    %19814 = stablehlo.broadcast_in_dim %19813, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19815 = stablehlo.broadcast_in_dim %19814, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19816 = stablehlo.divide %17424, %19815 : tensor<1x2048xf32> loc(#loc5272)
    %19817 = stablehlo.multiply %17432, %17432 : tensor<1x2048xf32> loc(#loc5268)
    %19818 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19819 = stablehlo.reduce(%19817 init: %19818) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19820 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19821 = stablehlo.broadcast_in_dim %19820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19822 = stablehlo.divide %19819, %19821 : tensor<1xf32> loc(#loc5270)
    %19823 = stablehlo.sqrt %19822 : tensor<1xf32> loc(#loc5249)
    %19824 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19825 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19826 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19827 = call @nan_to_num_2488(%19823, %19824, %19825, %19826) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19828 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19829 = stablehlo.broadcast_in_dim %19828, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19830 = stablehlo.divide %19827, %19829 : tensor<1xf32> loc(#loc5253)
    %19831 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19832 = stablehlo.broadcast_in_dim %19831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19833 = stablehlo.maximum %19832, %19830 : tensor<1xf32> loc(#loc5271)
    %19834 = stablehlo.broadcast_in_dim %19833, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19835 = stablehlo.broadcast_in_dim %19834, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19836 = stablehlo.divide %17432, %19835 : tensor<1x2048xf32> loc(#loc5272)
    %19837 = stablehlo.multiply %17440, %17440 : tensor<1x2048xf32> loc(#loc5268)
    %19838 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19839 = stablehlo.reduce(%19837 init: %19838) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19840 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19841 = stablehlo.broadcast_in_dim %19840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19842 = stablehlo.divide %19839, %19841 : tensor<1xf32> loc(#loc5270)
    %19843 = stablehlo.sqrt %19842 : tensor<1xf32> loc(#loc5249)
    %19844 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19845 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19846 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19847 = call @nan_to_num_2493(%19843, %19844, %19845, %19846) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19848 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19849 = stablehlo.broadcast_in_dim %19848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19850 = stablehlo.divide %19847, %19849 : tensor<1xf32> loc(#loc5253)
    %19851 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19852 = stablehlo.broadcast_in_dim %19851, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19853 = stablehlo.maximum %19852, %19850 : tensor<1xf32> loc(#loc5271)
    %19854 = stablehlo.broadcast_in_dim %19853, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19855 = stablehlo.broadcast_in_dim %19854, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19856 = stablehlo.divide %17440, %19855 : tensor<1x2048xf32> loc(#loc5272)
    %19857 = stablehlo.multiply %17448, %17448 : tensor<1x2048xf32> loc(#loc5268)
    %19858 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19859 = stablehlo.reduce(%19857 init: %19858) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19860 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19861 = stablehlo.broadcast_in_dim %19860, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19862 = stablehlo.divide %19859, %19861 : tensor<1xf32> loc(#loc5270)
    %19863 = stablehlo.sqrt %19862 : tensor<1xf32> loc(#loc5249)
    %19864 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19865 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19866 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19867 = call @nan_to_num_2498(%19863, %19864, %19865, %19866) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19868 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19869 = stablehlo.broadcast_in_dim %19868, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19870 = stablehlo.divide %19867, %19869 : tensor<1xf32> loc(#loc5253)
    %19871 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19872 = stablehlo.broadcast_in_dim %19871, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19873 = stablehlo.maximum %19872, %19870 : tensor<1xf32> loc(#loc5271)
    %19874 = stablehlo.broadcast_in_dim %19873, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19875 = stablehlo.broadcast_in_dim %19874, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19876 = stablehlo.divide %17448, %19875 : tensor<1x2048xf32> loc(#loc5272)
    %19877 = stablehlo.multiply %17456, %17456 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %19878 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %19879 = stablehlo.reduce(%19877 init: %19878) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %19880 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %19881 = stablehlo.broadcast_in_dim %19880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19882 = stablehlo.divide %19879, %19881 : tensor<1xf32> loc(#loc5270)
    %19883 = stablehlo.sqrt %19882 : tensor<1xf32> loc(#loc5249)
    %19884 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19885 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19886 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19887 = call @nan_to_num_2503(%19883, %19884, %19885, %19886) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19888 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19889 = stablehlo.broadcast_in_dim %19888, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19890 = stablehlo.divide %19887, %19889 : tensor<1xf32> loc(#loc5253)
    %19891 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19892 = stablehlo.broadcast_in_dim %19891, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19893 = stablehlo.maximum %19892, %19890 : tensor<1xf32> loc(#loc5271)
    %19894 = stablehlo.broadcast_in_dim %19893, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %19895 = stablehlo.broadcast_in_dim %19894, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19896 = stablehlo.divide %17456, %19895 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %19897 = stablehlo.multiply %17464, %17464 : tensor<1x64xf32> loc(#loc5268)
    %19898 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19899 = stablehlo.reduce(%19897 init: %19898) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19900 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %19901 = stablehlo.broadcast_in_dim %19900, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19902 = stablehlo.divide %19899, %19901 : tensor<1xf32> loc(#loc5270)
    %19903 = stablehlo.sqrt %19902 : tensor<1xf32> loc(#loc5249)
    %19904 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19905 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19906 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19907 = call @nan_to_num_2508(%19903, %19904, %19905, %19906) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19908 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19909 = stablehlo.broadcast_in_dim %19908, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19910 = stablehlo.divide %19907, %19909 : tensor<1xf32> loc(#loc5253)
    %19911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19912 = stablehlo.broadcast_in_dim %19911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19913 = stablehlo.maximum %19912, %19910 : tensor<1xf32> loc(#loc5271)
    %19914 = stablehlo.broadcast_in_dim %19913, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19915 = stablehlo.broadcast_in_dim %19914, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %19916 = stablehlo.divide %17464, %19915 : tensor<1x64xf32> loc(#loc5272)
    %19917 = stablehlo.multiply %17472, %17472 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %19918 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %19919 = stablehlo.reduce(%19917 init: %19918) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %19920 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %19921 = stablehlo.broadcast_in_dim %19920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19922 = stablehlo.divide %19919, %19921 : tensor<1xf32> loc(#loc5270)
    %19923 = stablehlo.sqrt %19922 : tensor<1xf32> loc(#loc5249)
    %19924 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19925 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19926 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19927 = call @nan_to_num_2513(%19923, %19924, %19925, %19926) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19928 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19929 = stablehlo.broadcast_in_dim %19928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19930 = stablehlo.divide %19927, %19929 : tensor<1xf32> loc(#loc5253)
    %19931 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19932 = stablehlo.broadcast_in_dim %19931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19933 = stablehlo.maximum %19932, %19930 : tensor<1xf32> loc(#loc5271)
    %19934 = stablehlo.broadcast_in_dim %19933, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %19935 = stablehlo.broadcast_in_dim %19934, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %19936 = stablehlo.divide %17472, %19935 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %19937 = stablehlo.multiply %17480, %17480 : tensor<1x8192xf32> loc(#loc5268)
    %19938 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19939 = stablehlo.reduce(%19937 init: %19938) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19940 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %19941 = stablehlo.broadcast_in_dim %19940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19942 = stablehlo.divide %19939, %19941 : tensor<1xf32> loc(#loc5270)
    %19943 = stablehlo.sqrt %19942 : tensor<1xf32> loc(#loc5249)
    %19944 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19945 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19946 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19947 = call @nan_to_num_2518(%19943, %19944, %19945, %19946) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19948 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19949 = stablehlo.broadcast_in_dim %19948, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19950 = stablehlo.divide %19947, %19949 : tensor<1xf32> loc(#loc5253)
    %19951 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19952 = stablehlo.broadcast_in_dim %19951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19953 = stablehlo.maximum %19952, %19950 : tensor<1xf32> loc(#loc5271)
    %19954 = stablehlo.broadcast_in_dim %19953, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19955 = stablehlo.broadcast_in_dim %19954, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %19956 = stablehlo.divide %17480, %19955 : tensor<1x8192xf32> loc(#loc5272)
    %19957 = stablehlo.multiply %17488, %17488 : tensor<1x2048x8192xf32> loc(#loc5268)
    %19958 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19959 = stablehlo.reduce(%19957 init: %19958) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %19960 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %19961 = stablehlo.broadcast_in_dim %19960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19962 = stablehlo.divide %19959, %19961 : tensor<1xf32> loc(#loc5270)
    %19963 = stablehlo.sqrt %19962 : tensor<1xf32> loc(#loc5249)
    %19964 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19965 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19966 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19967 = call @nan_to_num_2523(%19963, %19964, %19965, %19966) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19969 = stablehlo.broadcast_in_dim %19968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19970 = stablehlo.divide %19967, %19969 : tensor<1xf32> loc(#loc5253)
    %19971 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19972 = stablehlo.broadcast_in_dim %19971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19973 = stablehlo.maximum %19972, %19970 : tensor<1xf32> loc(#loc5271)
    %19974 = stablehlo.broadcast_in_dim %19973, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %19975 = stablehlo.broadcast_in_dim %19974, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %19976 = stablehlo.divide %17488, %19975 : tensor<1x2048x8192xf32> loc(#loc5272)
    %19977 = stablehlo.multiply %17496, %17496 : tensor<1x2048xf32> loc(#loc5268)
    %19978 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %19979 = stablehlo.reduce(%19977 init: %19978) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %19980 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %19981 = stablehlo.broadcast_in_dim %19980, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %19982 = stablehlo.divide %19979, %19981 : tensor<1xf32> loc(#loc5270)
    %19983 = stablehlo.sqrt %19982 : tensor<1xf32> loc(#loc5249)
    %19984 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19985 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19986 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %19987 = call @nan_to_num_2528(%19983, %19984, %19985, %19986) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %19988 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19989 = stablehlo.broadcast_in_dim %19988, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %19990 = stablehlo.divide %19987, %19989 : tensor<1xf32> loc(#loc5253)
    %19991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %19992 = stablehlo.broadcast_in_dim %19991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %19993 = stablehlo.maximum %19992, %19990 : tensor<1xf32> loc(#loc5271)
    %19994 = stablehlo.broadcast_in_dim %19993, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %19995 = stablehlo.broadcast_in_dim %19994, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %19996 = stablehlo.divide %17496, %19995 : tensor<1x2048xf32> loc(#loc5272)
    %19997 = stablehlo.multiply %17504, %17504 : tensor<1x8192x2048xf32> loc(#loc5268)
    %19998 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %19999 = stablehlo.reduce(%19997 init: %19998) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20000 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20001 = stablehlo.broadcast_in_dim %20000, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20002 = stablehlo.divide %19999, %20001 : tensor<1xf32> loc(#loc5270)
    %20003 = stablehlo.sqrt %20002 : tensor<1xf32> loc(#loc5249)
    %20004 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20005 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20006 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20007 = call @nan_to_num_2533(%20003, %20004, %20005, %20006) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20008 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20009 = stablehlo.broadcast_in_dim %20008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20010 = stablehlo.divide %20007, %20009 : tensor<1xf32> loc(#loc5253)
    %20011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20012 = stablehlo.broadcast_in_dim %20011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20013 = stablehlo.maximum %20012, %20010 : tensor<1xf32> loc(#loc5271)
    %20014 = stablehlo.broadcast_in_dim %20013, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20015 = stablehlo.broadcast_in_dim %20014, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %20016 = stablehlo.divide %17504, %20015 : tensor<1x8192x2048xf32> loc(#loc5272)
    %20017 = stablehlo.multiply %17512, %17512 : tensor<1x2048xf32> loc(#loc5268)
    %20018 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20019 = stablehlo.reduce(%20017 init: %20018) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20020 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20021 = stablehlo.broadcast_in_dim %20020, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20022 = stablehlo.divide %20019, %20021 : tensor<1xf32> loc(#loc5270)
    %20023 = stablehlo.sqrt %20022 : tensor<1xf32> loc(#loc5249)
    %20024 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20025 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20026 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20027 = call @nan_to_num_2538(%20023, %20024, %20025, %20026) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20028 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20029 = stablehlo.broadcast_in_dim %20028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20030 = stablehlo.divide %20027, %20029 : tensor<1xf32> loc(#loc5253)
    %20031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20032 = stablehlo.broadcast_in_dim %20031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20033 = stablehlo.maximum %20032, %20030 : tensor<1xf32> loc(#loc5271)
    %20034 = stablehlo.broadcast_in_dim %20033, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20035 = stablehlo.broadcast_in_dim %20034, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20036 = stablehlo.divide %17512, %20035 : tensor<1x2048xf32> loc(#loc5272)
    %20037 = stablehlo.multiply %17520, %17520 : tensor<1x2048xf32> loc(#loc5268)
    %20038 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20039 = stablehlo.reduce(%20037 init: %20038) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20040 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20041 = stablehlo.broadcast_in_dim %20040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20042 = stablehlo.divide %20039, %20041 : tensor<1xf32> loc(#loc5270)
    %20043 = stablehlo.sqrt %20042 : tensor<1xf32> loc(#loc5249)
    %20044 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20045 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20046 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20047 = call @nan_to_num_2543(%20043, %20044, %20045, %20046) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20048 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20049 = stablehlo.broadcast_in_dim %20048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20050 = stablehlo.divide %20047, %20049 : tensor<1xf32> loc(#loc5253)
    %20051 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20052 = stablehlo.broadcast_in_dim %20051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20053 = stablehlo.maximum %20052, %20050 : tensor<1xf32> loc(#loc5271)
    %20054 = stablehlo.broadcast_in_dim %20053, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20055 = stablehlo.broadcast_in_dim %20054, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20056 = stablehlo.divide %17520, %20055 : tensor<1x2048xf32> loc(#loc5272)
    %20057 = stablehlo.multiply %17528, %17528 : tensor<1x2048xf32> loc(#loc5268)
    %20058 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20059 = stablehlo.reduce(%20057 init: %20058) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20060 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20061 = stablehlo.broadcast_in_dim %20060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20062 = stablehlo.divide %20059, %20061 : tensor<1xf32> loc(#loc5270)
    %20063 = stablehlo.sqrt %20062 : tensor<1xf32> loc(#loc5249)
    %20064 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20065 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20066 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20067 = call @nan_to_num_2548(%20063, %20064, %20065, %20066) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20069 = stablehlo.broadcast_in_dim %20068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20070 = stablehlo.divide %20067, %20069 : tensor<1xf32> loc(#loc5253)
    %20071 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20072 = stablehlo.broadcast_in_dim %20071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20073 = stablehlo.maximum %20072, %20070 : tensor<1xf32> loc(#loc5271)
    %20074 = stablehlo.broadcast_in_dim %20073, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20075 = stablehlo.broadcast_in_dim %20074, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20076 = stablehlo.divide %17528, %20075 : tensor<1x2048xf32> loc(#loc5272)
    %20077 = stablehlo.multiply %17536, %17536 : tensor<1x2048xf32> loc(#loc5268)
    %20078 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20079 = stablehlo.reduce(%20077 init: %20078) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20080 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20081 = stablehlo.broadcast_in_dim %20080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20082 = stablehlo.divide %20079, %20081 : tensor<1xf32> loc(#loc5270)
    %20083 = stablehlo.sqrt %20082 : tensor<1xf32> loc(#loc5249)
    %20084 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20085 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20086 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20087 = call @nan_to_num_2553(%20083, %20084, %20085, %20086) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20089 = stablehlo.broadcast_in_dim %20088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20090 = stablehlo.divide %20087, %20089 : tensor<1xf32> loc(#loc5253)
    %20091 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20092 = stablehlo.broadcast_in_dim %20091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20093 = stablehlo.maximum %20092, %20090 : tensor<1xf32> loc(#loc5271)
    %20094 = stablehlo.broadcast_in_dim %20093, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20095 = stablehlo.broadcast_in_dim %20094, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20096 = stablehlo.divide %17536, %20095 : tensor<1x2048xf32> loc(#loc5272)
    %20097 = stablehlo.multiply %17544, %17544 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %20098 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %20099 = stablehlo.reduce(%20097 init: %20098) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %20100 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %20101 = stablehlo.broadcast_in_dim %20100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20102 = stablehlo.divide %20099, %20101 : tensor<1xf32> loc(#loc5270)
    %20103 = stablehlo.sqrt %20102 : tensor<1xf32> loc(#loc5249)
    %20104 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20105 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20106 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20107 = call @nan_to_num_2558(%20103, %20104, %20105, %20106) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20108 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20109 = stablehlo.broadcast_in_dim %20108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20110 = stablehlo.divide %20107, %20109 : tensor<1xf32> loc(#loc5253)
    %20111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20112 = stablehlo.broadcast_in_dim %20111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20113 = stablehlo.maximum %20112, %20110 : tensor<1xf32> loc(#loc5271)
    %20114 = stablehlo.broadcast_in_dim %20113, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %20115 = stablehlo.broadcast_in_dim %20114, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20116 = stablehlo.divide %17544, %20115 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20117 = stablehlo.multiply %17552, %17552 : tensor<1x64xf32> loc(#loc5268)
    %20118 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20119 = stablehlo.reduce(%20117 init: %20118) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20120 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %20121 = stablehlo.broadcast_in_dim %20120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20122 = stablehlo.divide %20119, %20121 : tensor<1xf32> loc(#loc5270)
    %20123 = stablehlo.sqrt %20122 : tensor<1xf32> loc(#loc5249)
    %20124 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20125 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20126 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20127 = call @nan_to_num_2563(%20123, %20124, %20125, %20126) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20128 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20129 = stablehlo.broadcast_in_dim %20128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20130 = stablehlo.divide %20127, %20129 : tensor<1xf32> loc(#loc5253)
    %20131 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20132 = stablehlo.broadcast_in_dim %20131, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20133 = stablehlo.maximum %20132, %20130 : tensor<1xf32> loc(#loc5271)
    %20134 = stablehlo.broadcast_in_dim %20133, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20135 = stablehlo.broadcast_in_dim %20134, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %20136 = stablehlo.divide %17552, %20135 : tensor<1x64xf32> loc(#loc5272)
    %20137 = stablehlo.multiply %17560, %17560 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %20138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %20139 = stablehlo.reduce(%20137 init: %20138) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %20140 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %20141 = stablehlo.broadcast_in_dim %20140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20142 = stablehlo.divide %20139, %20141 : tensor<1xf32> loc(#loc5270)
    %20143 = stablehlo.sqrt %20142 : tensor<1xf32> loc(#loc5249)
    %20144 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20145 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20146 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20147 = call @nan_to_num_2568(%20143, %20144, %20145, %20146) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20148 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20149 = stablehlo.broadcast_in_dim %20148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20150 = stablehlo.divide %20147, %20149 : tensor<1xf32> loc(#loc5253)
    %20151 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20152 = stablehlo.broadcast_in_dim %20151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20153 = stablehlo.maximum %20152, %20150 : tensor<1xf32> loc(#loc5271)
    %20154 = stablehlo.broadcast_in_dim %20153, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %20155 = stablehlo.broadcast_in_dim %20154, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %20156 = stablehlo.divide %17560, %20155 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %20157 = stablehlo.multiply %17568, %17568 : tensor<1x8192xf32> loc(#loc5268)
    %20158 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20159 = stablehlo.reduce(%20157 init: %20158) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20160 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %20161 = stablehlo.broadcast_in_dim %20160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20162 = stablehlo.divide %20159, %20161 : tensor<1xf32> loc(#loc5270)
    %20163 = stablehlo.sqrt %20162 : tensor<1xf32> loc(#loc5249)
    %20164 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20165 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20166 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20167 = call @nan_to_num_2573(%20163, %20164, %20165, %20166) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20168 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20169 = stablehlo.broadcast_in_dim %20168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20170 = stablehlo.divide %20167, %20169 : tensor<1xf32> loc(#loc5253)
    %20171 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20172 = stablehlo.broadcast_in_dim %20171, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20173 = stablehlo.maximum %20172, %20170 : tensor<1xf32> loc(#loc5271)
    %20174 = stablehlo.broadcast_in_dim %20173, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20175 = stablehlo.broadcast_in_dim %20174, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %20176 = stablehlo.divide %17568, %20175 : tensor<1x8192xf32> loc(#loc5272)
    %20177 = stablehlo.multiply %17576, %17576 : tensor<1x2048x8192xf32> loc(#loc5268)
    %20178 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20179 = stablehlo.reduce(%20177 init: %20178) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20180 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20181 = stablehlo.broadcast_in_dim %20180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20182 = stablehlo.divide %20179, %20181 : tensor<1xf32> loc(#loc5270)
    %20183 = stablehlo.sqrt %20182 : tensor<1xf32> loc(#loc5249)
    %20184 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20185 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20186 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20187 = call @nan_to_num_2578(%20183, %20184, %20185, %20186) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20189 = stablehlo.broadcast_in_dim %20188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20190 = stablehlo.divide %20187, %20189 : tensor<1xf32> loc(#loc5253)
    %20191 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20192 = stablehlo.broadcast_in_dim %20191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20193 = stablehlo.maximum %20192, %20190 : tensor<1xf32> loc(#loc5271)
    %20194 = stablehlo.broadcast_in_dim %20193, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20195 = stablehlo.broadcast_in_dim %20194, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %20196 = stablehlo.divide %17576, %20195 : tensor<1x2048x8192xf32> loc(#loc5272)
    %20197 = stablehlo.multiply %17584, %17584 : tensor<1x2048xf32> loc(#loc5268)
    %20198 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20199 = stablehlo.reduce(%20197 init: %20198) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20200 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20201 = stablehlo.broadcast_in_dim %20200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20202 = stablehlo.divide %20199, %20201 : tensor<1xf32> loc(#loc5270)
    %20203 = stablehlo.sqrt %20202 : tensor<1xf32> loc(#loc5249)
    %20204 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20205 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20206 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20207 = call @nan_to_num_2583(%20203, %20204, %20205, %20206) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20209 = stablehlo.broadcast_in_dim %20208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20210 = stablehlo.divide %20207, %20209 : tensor<1xf32> loc(#loc5253)
    %20211 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20212 = stablehlo.broadcast_in_dim %20211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20213 = stablehlo.maximum %20212, %20210 : tensor<1xf32> loc(#loc5271)
    %20214 = stablehlo.broadcast_in_dim %20213, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20215 = stablehlo.broadcast_in_dim %20214, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20216 = stablehlo.divide %17584, %20215 : tensor<1x2048xf32> loc(#loc5272)
    %20217 = stablehlo.multiply %17592, %17592 : tensor<1x8192x2048xf32> loc(#loc5268)
    %20218 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20219 = stablehlo.reduce(%20217 init: %20218) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20220 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20221 = stablehlo.broadcast_in_dim %20220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20222 = stablehlo.divide %20219, %20221 : tensor<1xf32> loc(#loc5270)
    %20223 = stablehlo.sqrt %20222 : tensor<1xf32> loc(#loc5249)
    %20224 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20225 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20226 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20227 = call @nan_to_num_2588(%20223, %20224, %20225, %20226) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20228 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20229 = stablehlo.broadcast_in_dim %20228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20230 = stablehlo.divide %20227, %20229 : tensor<1xf32> loc(#loc5253)
    %20231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20232 = stablehlo.broadcast_in_dim %20231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20233 = stablehlo.maximum %20232, %20230 : tensor<1xf32> loc(#loc5271)
    %20234 = stablehlo.broadcast_in_dim %20233, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20235 = stablehlo.broadcast_in_dim %20234, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %20236 = stablehlo.divide %17592, %20235 : tensor<1x8192x2048xf32> loc(#loc5272)
    %20237 = stablehlo.multiply %17600, %17600 : tensor<1x2048xf32> loc(#loc5268)
    %20238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20239 = stablehlo.reduce(%20237 init: %20238) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20240 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20241 = stablehlo.broadcast_in_dim %20240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20242 = stablehlo.divide %20239, %20241 : tensor<1xf32> loc(#loc5270)
    %20243 = stablehlo.sqrt %20242 : tensor<1xf32> loc(#loc5249)
    %20244 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20245 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20246 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20247 = call @nan_to_num_2593(%20243, %20244, %20245, %20246) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20248 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20249 = stablehlo.broadcast_in_dim %20248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20250 = stablehlo.divide %20247, %20249 : tensor<1xf32> loc(#loc5253)
    %20251 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20252 = stablehlo.broadcast_in_dim %20251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20253 = stablehlo.maximum %20252, %20250 : tensor<1xf32> loc(#loc5271)
    %20254 = stablehlo.broadcast_in_dim %20253, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20255 = stablehlo.broadcast_in_dim %20254, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20256 = stablehlo.divide %17600, %20255 : tensor<1x2048xf32> loc(#loc5272)
    %20257 = stablehlo.multiply %17608, %17608 : tensor<1x2048xf32> loc(#loc5268)
    %20258 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20259 = stablehlo.reduce(%20257 init: %20258) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20260 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20261 = stablehlo.broadcast_in_dim %20260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20262 = stablehlo.divide %20259, %20261 : tensor<1xf32> loc(#loc5270)
    %20263 = stablehlo.sqrt %20262 : tensor<1xf32> loc(#loc5249)
    %20264 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20265 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20266 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20267 = call @nan_to_num_2598(%20263, %20264, %20265, %20266) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20268 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20269 = stablehlo.broadcast_in_dim %20268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20270 = stablehlo.divide %20267, %20269 : tensor<1xf32> loc(#loc5253)
    %20271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20272 = stablehlo.broadcast_in_dim %20271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20273 = stablehlo.maximum %20272, %20270 : tensor<1xf32> loc(#loc5271)
    %20274 = stablehlo.broadcast_in_dim %20273, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20275 = stablehlo.broadcast_in_dim %20274, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20276 = stablehlo.divide %17608, %20275 : tensor<1x2048xf32> loc(#loc5272)
    %20277 = stablehlo.multiply %17616, %17616 : tensor<1x2048xf32> loc(#loc5268)
    %20278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20279 = stablehlo.reduce(%20277 init: %20278) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20280 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20281 = stablehlo.broadcast_in_dim %20280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20282 = stablehlo.divide %20279, %20281 : tensor<1xf32> loc(#loc5270)
    %20283 = stablehlo.sqrt %20282 : tensor<1xf32> loc(#loc5249)
    %20284 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20285 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20286 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20287 = call @nan_to_num_2603(%20283, %20284, %20285, %20286) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20288 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20289 = stablehlo.broadcast_in_dim %20288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20290 = stablehlo.divide %20287, %20289 : tensor<1xf32> loc(#loc5253)
    %20291 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20292 = stablehlo.broadcast_in_dim %20291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20293 = stablehlo.maximum %20292, %20290 : tensor<1xf32> loc(#loc5271)
    %20294 = stablehlo.broadcast_in_dim %20293, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20295 = stablehlo.broadcast_in_dim %20294, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20296 = stablehlo.divide %17616, %20295 : tensor<1x2048xf32> loc(#loc5272)
    %20297 = stablehlo.multiply %17624, %17624 : tensor<1x2048xf32> loc(#loc5268)
    %20298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20299 = stablehlo.reduce(%20297 init: %20298) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20300 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20301 = stablehlo.broadcast_in_dim %20300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20302 = stablehlo.divide %20299, %20301 : tensor<1xf32> loc(#loc5270)
    %20303 = stablehlo.sqrt %20302 : tensor<1xf32> loc(#loc5249)
    %20304 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20305 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20306 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20307 = call @nan_to_num_2608(%20303, %20304, %20305, %20306) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20308 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20309 = stablehlo.broadcast_in_dim %20308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20310 = stablehlo.divide %20307, %20309 : tensor<1xf32> loc(#loc5253)
    %20311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20312 = stablehlo.broadcast_in_dim %20311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20313 = stablehlo.maximum %20312, %20310 : tensor<1xf32> loc(#loc5271)
    %20314 = stablehlo.broadcast_in_dim %20313, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20315 = stablehlo.broadcast_in_dim %20314, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20316 = stablehlo.divide %17624, %20315 : tensor<1x2048xf32> loc(#loc5272)
    %20317 = stablehlo.multiply %17632, %17632 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %20318 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %20319 = stablehlo.reduce(%20317 init: %20318) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %20320 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %20321 = stablehlo.broadcast_in_dim %20320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20322 = stablehlo.divide %20319, %20321 : tensor<1xf32> loc(#loc5270)
    %20323 = stablehlo.sqrt %20322 : tensor<1xf32> loc(#loc5249)
    %20324 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20325 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20326 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20327 = call @nan_to_num_2613(%20323, %20324, %20325, %20326) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20328 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20329 = stablehlo.broadcast_in_dim %20328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20330 = stablehlo.divide %20327, %20329 : tensor<1xf32> loc(#loc5253)
    %20331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20332 = stablehlo.broadcast_in_dim %20331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20333 = stablehlo.maximum %20332, %20330 : tensor<1xf32> loc(#loc5271)
    %20334 = stablehlo.broadcast_in_dim %20333, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %20335 = stablehlo.broadcast_in_dim %20334, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20336 = stablehlo.divide %17632, %20335 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20337 = stablehlo.multiply %17640, %17640 : tensor<1x64xf32> loc(#loc5268)
    %20338 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20339 = stablehlo.reduce(%20337 init: %20338) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20340 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %20341 = stablehlo.broadcast_in_dim %20340, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20342 = stablehlo.divide %20339, %20341 : tensor<1xf32> loc(#loc5270)
    %20343 = stablehlo.sqrt %20342 : tensor<1xf32> loc(#loc5249)
    %20344 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20345 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20346 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20347 = call @nan_to_num_2618(%20343, %20344, %20345, %20346) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20348 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20349 = stablehlo.broadcast_in_dim %20348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20350 = stablehlo.divide %20347, %20349 : tensor<1xf32> loc(#loc5253)
    %20351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20352 = stablehlo.broadcast_in_dim %20351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20353 = stablehlo.maximum %20352, %20350 : tensor<1xf32> loc(#loc5271)
    %20354 = stablehlo.broadcast_in_dim %20353, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20355 = stablehlo.broadcast_in_dim %20354, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %20356 = stablehlo.divide %17640, %20355 : tensor<1x64xf32> loc(#loc5272)
    %20357 = stablehlo.multiply %17648, %17648 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %20358 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %20359 = stablehlo.reduce(%20357 init: %20358) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %20360 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %20361 = stablehlo.broadcast_in_dim %20360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20362 = stablehlo.divide %20359, %20361 : tensor<1xf32> loc(#loc5270)
    %20363 = stablehlo.sqrt %20362 : tensor<1xf32> loc(#loc5249)
    %20364 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20365 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20366 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20367 = call @nan_to_num_2623(%20363, %20364, %20365, %20366) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20369 = stablehlo.broadcast_in_dim %20368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20370 = stablehlo.divide %20367, %20369 : tensor<1xf32> loc(#loc5253)
    %20371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20372 = stablehlo.broadcast_in_dim %20371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20373 = stablehlo.maximum %20372, %20370 : tensor<1xf32> loc(#loc5271)
    %20374 = stablehlo.broadcast_in_dim %20373, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %20375 = stablehlo.broadcast_in_dim %20374, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %20376 = stablehlo.divide %17648, %20375 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %20377 = stablehlo.multiply %17656, %17656 : tensor<1x8192xf32> loc(#loc5268)
    %20378 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20379 = stablehlo.reduce(%20377 init: %20378) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20380 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %20381 = stablehlo.broadcast_in_dim %20380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20382 = stablehlo.divide %20379, %20381 : tensor<1xf32> loc(#loc5270)
    %20383 = stablehlo.sqrt %20382 : tensor<1xf32> loc(#loc5249)
    %20384 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20385 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20386 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20387 = call @nan_to_num_2628(%20383, %20384, %20385, %20386) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20389 = stablehlo.broadcast_in_dim %20388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20390 = stablehlo.divide %20387, %20389 : tensor<1xf32> loc(#loc5253)
    %20391 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20392 = stablehlo.broadcast_in_dim %20391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20393 = stablehlo.maximum %20392, %20390 : tensor<1xf32> loc(#loc5271)
    %20394 = stablehlo.broadcast_in_dim %20393, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20395 = stablehlo.broadcast_in_dim %20394, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %20396 = stablehlo.divide %17656, %20395 : tensor<1x8192xf32> loc(#loc5272)
    %20397 = stablehlo.multiply %17664, %17664 : tensor<1x2048x8192xf32> loc(#loc5268)
    %20398 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20399 = stablehlo.reduce(%20397 init: %20398) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20400 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20401 = stablehlo.broadcast_in_dim %20400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20402 = stablehlo.divide %20399, %20401 : tensor<1xf32> loc(#loc5270)
    %20403 = stablehlo.sqrt %20402 : tensor<1xf32> loc(#loc5249)
    %20404 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20405 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20406 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20407 = call @nan_to_num_2633(%20403, %20404, %20405, %20406) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20409 = stablehlo.broadcast_in_dim %20408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20410 = stablehlo.divide %20407, %20409 : tensor<1xf32> loc(#loc5253)
    %20411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20412 = stablehlo.broadcast_in_dim %20411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20413 = stablehlo.maximum %20412, %20410 : tensor<1xf32> loc(#loc5271)
    %20414 = stablehlo.broadcast_in_dim %20413, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20415 = stablehlo.broadcast_in_dim %20414, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %20416 = stablehlo.divide %17664, %20415 : tensor<1x2048x8192xf32> loc(#loc5272)
    %20417 = stablehlo.multiply %17672, %17672 : tensor<1x2048xf32> loc(#loc5268)
    %20418 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20419 = stablehlo.reduce(%20417 init: %20418) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20420 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20421 = stablehlo.broadcast_in_dim %20420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20422 = stablehlo.divide %20419, %20421 : tensor<1xf32> loc(#loc5270)
    %20423 = stablehlo.sqrt %20422 : tensor<1xf32> loc(#loc5249)
    %20424 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20425 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20426 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20427 = call @nan_to_num_2638(%20423, %20424, %20425, %20426) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20429 = stablehlo.broadcast_in_dim %20428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20430 = stablehlo.divide %20427, %20429 : tensor<1xf32> loc(#loc5253)
    %20431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20432 = stablehlo.broadcast_in_dim %20431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20433 = stablehlo.maximum %20432, %20430 : tensor<1xf32> loc(#loc5271)
    %20434 = stablehlo.broadcast_in_dim %20433, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20435 = stablehlo.broadcast_in_dim %20434, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20436 = stablehlo.divide %17672, %20435 : tensor<1x2048xf32> loc(#loc5272)
    %20437 = stablehlo.multiply %17680, %17680 : tensor<1x8192x2048xf32> loc(#loc5268)
    %20438 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20439 = stablehlo.reduce(%20437 init: %20438) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20440 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20441 = stablehlo.broadcast_in_dim %20440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20442 = stablehlo.divide %20439, %20441 : tensor<1xf32> loc(#loc5270)
    %20443 = stablehlo.sqrt %20442 : tensor<1xf32> loc(#loc5249)
    %20444 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20445 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20446 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20447 = call @nan_to_num_2643(%20443, %20444, %20445, %20446) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20448 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20449 = stablehlo.broadcast_in_dim %20448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20450 = stablehlo.divide %20447, %20449 : tensor<1xf32> loc(#loc5253)
    %20451 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20452 = stablehlo.broadcast_in_dim %20451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20453 = stablehlo.maximum %20452, %20450 : tensor<1xf32> loc(#loc5271)
    %20454 = stablehlo.broadcast_in_dim %20453, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20455 = stablehlo.broadcast_in_dim %20454, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %20456 = stablehlo.divide %17680, %20455 : tensor<1x8192x2048xf32> loc(#loc5272)
    %20457 = stablehlo.multiply %17688, %17688 : tensor<1x2048xf32> loc(#loc5268)
    %20458 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20459 = stablehlo.reduce(%20457 init: %20458) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20460 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20461 = stablehlo.broadcast_in_dim %20460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20462 = stablehlo.divide %20459, %20461 : tensor<1xf32> loc(#loc5270)
    %20463 = stablehlo.sqrt %20462 : tensor<1xf32> loc(#loc5249)
    %20464 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20465 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20466 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20467 = call @nan_to_num_2648(%20463, %20464, %20465, %20466) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20468 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20469 = stablehlo.broadcast_in_dim %20468, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20470 = stablehlo.divide %20467, %20469 : tensor<1xf32> loc(#loc5253)
    %20471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20472 = stablehlo.broadcast_in_dim %20471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20473 = stablehlo.maximum %20472, %20470 : tensor<1xf32> loc(#loc5271)
    %20474 = stablehlo.broadcast_in_dim %20473, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20475 = stablehlo.broadcast_in_dim %20474, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20476 = stablehlo.divide %17688, %20475 : tensor<1x2048xf32> loc(#loc5272)
    %20477 = stablehlo.multiply %17696, %17696 : tensor<1x2048xf32> loc(#loc5268)
    %20478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20479 = stablehlo.reduce(%20477 init: %20478) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20480 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20481 = stablehlo.broadcast_in_dim %20480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20482 = stablehlo.divide %20479, %20481 : tensor<1xf32> loc(#loc5270)
    %20483 = stablehlo.sqrt %20482 : tensor<1xf32> loc(#loc5249)
    %20484 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20485 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20486 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20487 = call @nan_to_num_2653(%20483, %20484, %20485, %20486) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20488 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20489 = stablehlo.broadcast_in_dim %20488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20490 = stablehlo.divide %20487, %20489 : tensor<1xf32> loc(#loc5253)
    %20491 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20492 = stablehlo.broadcast_in_dim %20491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20493 = stablehlo.maximum %20492, %20490 : tensor<1xf32> loc(#loc5271)
    %20494 = stablehlo.broadcast_in_dim %20493, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20495 = stablehlo.broadcast_in_dim %20494, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20496 = stablehlo.divide %17696, %20495 : tensor<1x2048xf32> loc(#loc5272)
    %20497 = stablehlo.multiply %17704, %17704 : tensor<1x2048xf32> loc(#loc5268)
    %20498 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20499 = stablehlo.reduce(%20497 init: %20498) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20500 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20501 = stablehlo.broadcast_in_dim %20500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20502 = stablehlo.divide %20499, %20501 : tensor<1xf32> loc(#loc5270)
    %20503 = stablehlo.sqrt %20502 : tensor<1xf32> loc(#loc5249)
    %20504 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20505 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20506 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20507 = call @nan_to_num_2658(%20503, %20504, %20505, %20506) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20508 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20509 = stablehlo.broadcast_in_dim %20508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20510 = stablehlo.divide %20507, %20509 : tensor<1xf32> loc(#loc5253)
    %20511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20512 = stablehlo.broadcast_in_dim %20511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20513 = stablehlo.maximum %20512, %20510 : tensor<1xf32> loc(#loc5271)
    %20514 = stablehlo.broadcast_in_dim %20513, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20515 = stablehlo.broadcast_in_dim %20514, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20516 = stablehlo.divide %17704, %20515 : tensor<1x2048xf32> loc(#loc5272)
    %20517 = stablehlo.multiply %17712, %17712 : tensor<1x2048xf32> loc(#loc5268)
    %20518 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20519 = stablehlo.reduce(%20517 init: %20518) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20520 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20521 = stablehlo.broadcast_in_dim %20520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20522 = stablehlo.divide %20519, %20521 : tensor<1xf32> loc(#loc5270)
    %20523 = stablehlo.sqrt %20522 : tensor<1xf32> loc(#loc5249)
    %20524 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20525 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20526 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20527 = call @nan_to_num_2663(%20523, %20524, %20525, %20526) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20529 = stablehlo.broadcast_in_dim %20528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20530 = stablehlo.divide %20527, %20529 : tensor<1xf32> loc(#loc5253)
    %20531 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20532 = stablehlo.broadcast_in_dim %20531, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20533 = stablehlo.maximum %20532, %20530 : tensor<1xf32> loc(#loc5271)
    %20534 = stablehlo.broadcast_in_dim %20533, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20535 = stablehlo.broadcast_in_dim %20534, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20536 = stablehlo.divide %17712, %20535 : tensor<1x2048xf32> loc(#loc5272)
    %20537 = stablehlo.multiply %17720, %17720 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %20538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %20539 = stablehlo.reduce(%20537 init: %20538) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %20540 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %20541 = stablehlo.broadcast_in_dim %20540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20542 = stablehlo.divide %20539, %20541 : tensor<1xf32> loc(#loc5270)
    %20543 = stablehlo.sqrt %20542 : tensor<1xf32> loc(#loc5249)
    %20544 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20545 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20546 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20547 = call @nan_to_num_2668(%20543, %20544, %20545, %20546) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20548 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20549 = stablehlo.broadcast_in_dim %20548, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20550 = stablehlo.divide %20547, %20549 : tensor<1xf32> loc(#loc5253)
    %20551 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20552 = stablehlo.broadcast_in_dim %20551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20553 = stablehlo.maximum %20552, %20550 : tensor<1xf32> loc(#loc5271)
    %20554 = stablehlo.broadcast_in_dim %20553, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %20555 = stablehlo.broadcast_in_dim %20554, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20556 = stablehlo.divide %17720, %20555 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20557 = stablehlo.multiply %17728, %17728 : tensor<1x64xf32> loc(#loc5268)
    %20558 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20559 = stablehlo.reduce(%20557 init: %20558) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20560 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %20561 = stablehlo.broadcast_in_dim %20560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20562 = stablehlo.divide %20559, %20561 : tensor<1xf32> loc(#loc5270)
    %20563 = stablehlo.sqrt %20562 : tensor<1xf32> loc(#loc5249)
    %20564 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20565 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20566 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20567 = call @nan_to_num_2673(%20563, %20564, %20565, %20566) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20568 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20569 = stablehlo.broadcast_in_dim %20568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20570 = stablehlo.divide %20567, %20569 : tensor<1xf32> loc(#loc5253)
    %20571 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20572 = stablehlo.broadcast_in_dim %20571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20573 = stablehlo.maximum %20572, %20570 : tensor<1xf32> loc(#loc5271)
    %20574 = stablehlo.broadcast_in_dim %20573, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20575 = stablehlo.broadcast_in_dim %20574, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %20576 = stablehlo.divide %17728, %20575 : tensor<1x64xf32> loc(#loc5272)
    %20577 = stablehlo.multiply %17736, %17736 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %20578 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %20579 = stablehlo.reduce(%20577 init: %20578) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %20580 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %20581 = stablehlo.broadcast_in_dim %20580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20582 = stablehlo.divide %20579, %20581 : tensor<1xf32> loc(#loc5270)
    %20583 = stablehlo.sqrt %20582 : tensor<1xf32> loc(#loc5249)
    %20584 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20585 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20586 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20587 = call @nan_to_num_2678(%20583, %20584, %20585, %20586) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20588 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20589 = stablehlo.broadcast_in_dim %20588, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20590 = stablehlo.divide %20587, %20589 : tensor<1xf32> loc(#loc5253)
    %20591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20592 = stablehlo.broadcast_in_dim %20591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20593 = stablehlo.maximum %20592, %20590 : tensor<1xf32> loc(#loc5271)
    %20594 = stablehlo.broadcast_in_dim %20593, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %20595 = stablehlo.broadcast_in_dim %20594, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %20596 = stablehlo.divide %17736, %20595 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %20597 = stablehlo.multiply %17744, %17744 : tensor<1x8192xf32> loc(#loc5268)
    %20598 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20599 = stablehlo.reduce(%20597 init: %20598) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20600 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %20601 = stablehlo.broadcast_in_dim %20600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20602 = stablehlo.divide %20599, %20601 : tensor<1xf32> loc(#loc5270)
    %20603 = stablehlo.sqrt %20602 : tensor<1xf32> loc(#loc5249)
    %20604 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20605 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20606 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20607 = call @nan_to_num_2683(%20603, %20604, %20605, %20606) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20608 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20609 = stablehlo.broadcast_in_dim %20608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20610 = stablehlo.divide %20607, %20609 : tensor<1xf32> loc(#loc5253)
    %20611 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20612 = stablehlo.broadcast_in_dim %20611, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20613 = stablehlo.maximum %20612, %20610 : tensor<1xf32> loc(#loc5271)
    %20614 = stablehlo.broadcast_in_dim %20613, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20615 = stablehlo.broadcast_in_dim %20614, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %20616 = stablehlo.divide %17744, %20615 : tensor<1x8192xf32> loc(#loc5272)
    %20617 = stablehlo.multiply %17752, %17752 : tensor<1x2048x8192xf32> loc(#loc5268)
    %20618 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20619 = stablehlo.reduce(%20617 init: %20618) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20620 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20621 = stablehlo.broadcast_in_dim %20620, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20622 = stablehlo.divide %20619, %20621 : tensor<1xf32> loc(#loc5270)
    %20623 = stablehlo.sqrt %20622 : tensor<1xf32> loc(#loc5249)
    %20624 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20625 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20626 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20627 = call @nan_to_num_2688(%20623, %20624, %20625, %20626) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20628 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20629 = stablehlo.broadcast_in_dim %20628, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20630 = stablehlo.divide %20627, %20629 : tensor<1xf32> loc(#loc5253)
    %20631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20632 = stablehlo.broadcast_in_dim %20631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20633 = stablehlo.maximum %20632, %20630 : tensor<1xf32> loc(#loc5271)
    %20634 = stablehlo.broadcast_in_dim %20633, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20635 = stablehlo.broadcast_in_dim %20634, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %20636 = stablehlo.divide %17752, %20635 : tensor<1x2048x8192xf32> loc(#loc5272)
    %20637 = stablehlo.multiply %17760, %17760 : tensor<1x2048xf32> loc(#loc5268)
    %20638 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20639 = stablehlo.reduce(%20637 init: %20638) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20640 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20641 = stablehlo.broadcast_in_dim %20640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20642 = stablehlo.divide %20639, %20641 : tensor<1xf32> loc(#loc5270)
    %20643 = stablehlo.sqrt %20642 : tensor<1xf32> loc(#loc5249)
    %20644 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20645 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20646 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20647 = call @nan_to_num_2693(%20643, %20644, %20645, %20646) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20648 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20649 = stablehlo.broadcast_in_dim %20648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20650 = stablehlo.divide %20647, %20649 : tensor<1xf32> loc(#loc5253)
    %20651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20652 = stablehlo.broadcast_in_dim %20651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20653 = stablehlo.maximum %20652, %20650 : tensor<1xf32> loc(#loc5271)
    %20654 = stablehlo.broadcast_in_dim %20653, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20655 = stablehlo.broadcast_in_dim %20654, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20656 = stablehlo.divide %17760, %20655 : tensor<1x2048xf32> loc(#loc5272)
    %20657 = stablehlo.multiply %17768, %17768 : tensor<1x8192x2048xf32> loc(#loc5268)
    %20658 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20659 = stablehlo.reduce(%20657 init: %20658) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20660 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20661 = stablehlo.broadcast_in_dim %20660, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20662 = stablehlo.divide %20659, %20661 : tensor<1xf32> loc(#loc5270)
    %20663 = stablehlo.sqrt %20662 : tensor<1xf32> loc(#loc5249)
    %20664 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20665 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20666 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20667 = call @nan_to_num_2698(%20663, %20664, %20665, %20666) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20668 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20669 = stablehlo.broadcast_in_dim %20668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20670 = stablehlo.divide %20667, %20669 : tensor<1xf32> loc(#loc5253)
    %20671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20672 = stablehlo.broadcast_in_dim %20671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20673 = stablehlo.maximum %20672, %20670 : tensor<1xf32> loc(#loc5271)
    %20674 = stablehlo.broadcast_in_dim %20673, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20675 = stablehlo.broadcast_in_dim %20674, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %20676 = stablehlo.divide %17768, %20675 : tensor<1x8192x2048xf32> loc(#loc5272)
    %20677 = stablehlo.multiply %17776, %17776 : tensor<1x2048xf32> loc(#loc5268)
    %20678 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20679 = stablehlo.reduce(%20677 init: %20678) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20680 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20681 = stablehlo.broadcast_in_dim %20680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20682 = stablehlo.divide %20679, %20681 : tensor<1xf32> loc(#loc5270)
    %20683 = stablehlo.sqrt %20682 : tensor<1xf32> loc(#loc5249)
    %20684 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20685 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20686 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20687 = call @nan_to_num_2703(%20683, %20684, %20685, %20686) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20688 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20689 = stablehlo.broadcast_in_dim %20688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20690 = stablehlo.divide %20687, %20689 : tensor<1xf32> loc(#loc5253)
    %20691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20692 = stablehlo.broadcast_in_dim %20691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20693 = stablehlo.maximum %20692, %20690 : tensor<1xf32> loc(#loc5271)
    %20694 = stablehlo.broadcast_in_dim %20693, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20695 = stablehlo.broadcast_in_dim %20694, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20696 = stablehlo.divide %17776, %20695 : tensor<1x2048xf32> loc(#loc5272)
    %20697 = stablehlo.multiply %17784, %17784 : tensor<1x2048xf32> loc(#loc5268)
    %20698 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20699 = stablehlo.reduce(%20697 init: %20698) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20700 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20701 = stablehlo.broadcast_in_dim %20700, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20702 = stablehlo.divide %20699, %20701 : tensor<1xf32> loc(#loc5270)
    %20703 = stablehlo.sqrt %20702 : tensor<1xf32> loc(#loc5249)
    %20704 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20705 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20706 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20707 = call @nan_to_num_2708(%20703, %20704, %20705, %20706) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20708 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20709 = stablehlo.broadcast_in_dim %20708, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20710 = stablehlo.divide %20707, %20709 : tensor<1xf32> loc(#loc5253)
    %20711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20712 = stablehlo.broadcast_in_dim %20711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20713 = stablehlo.maximum %20712, %20710 : tensor<1xf32> loc(#loc5271)
    %20714 = stablehlo.broadcast_in_dim %20713, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20715 = stablehlo.broadcast_in_dim %20714, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20716 = stablehlo.divide %17784, %20715 : tensor<1x2048xf32> loc(#loc5272)
    %20717 = stablehlo.multiply %17792, %17792 : tensor<1x2048xf32> loc(#loc5268)
    %20718 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20719 = stablehlo.reduce(%20717 init: %20718) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20720 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20721 = stablehlo.broadcast_in_dim %20720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20722 = stablehlo.divide %20719, %20721 : tensor<1xf32> loc(#loc5270)
    %20723 = stablehlo.sqrt %20722 : tensor<1xf32> loc(#loc5249)
    %20724 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20725 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20726 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20727 = call @nan_to_num_2713(%20723, %20724, %20725, %20726) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20728 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20729 = stablehlo.broadcast_in_dim %20728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20730 = stablehlo.divide %20727, %20729 : tensor<1xf32> loc(#loc5253)
    %20731 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20732 = stablehlo.broadcast_in_dim %20731, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20733 = stablehlo.maximum %20732, %20730 : tensor<1xf32> loc(#loc5271)
    %20734 = stablehlo.broadcast_in_dim %20733, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20735 = stablehlo.broadcast_in_dim %20734, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20736 = stablehlo.divide %17792, %20735 : tensor<1x2048xf32> loc(#loc5272)
    %20737 = stablehlo.multiply %17800, %17800 : tensor<1x2048xf32> loc(#loc5268)
    %20738 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20739 = stablehlo.reduce(%20737 init: %20738) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20740 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20741 = stablehlo.broadcast_in_dim %20740, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20742 = stablehlo.divide %20739, %20741 : tensor<1xf32> loc(#loc5270)
    %20743 = stablehlo.sqrt %20742 : tensor<1xf32> loc(#loc5249)
    %20744 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20745 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20746 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20747 = call @nan_to_num_2718(%20743, %20744, %20745, %20746) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20748 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20749 = stablehlo.broadcast_in_dim %20748, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20750 = stablehlo.divide %20747, %20749 : tensor<1xf32> loc(#loc5253)
    %20751 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20752 = stablehlo.broadcast_in_dim %20751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20753 = stablehlo.maximum %20752, %20750 : tensor<1xf32> loc(#loc5271)
    %20754 = stablehlo.broadcast_in_dim %20753, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20755 = stablehlo.broadcast_in_dim %20754, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20756 = stablehlo.divide %17800, %20755 : tensor<1x2048xf32> loc(#loc5272)
    %20757 = stablehlo.multiply %17808, %17808 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %20758 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %20759 = stablehlo.reduce(%20757 init: %20758) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %20760 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %20761 = stablehlo.broadcast_in_dim %20760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20762 = stablehlo.divide %20759, %20761 : tensor<1xf32> loc(#loc5270)
    %20763 = stablehlo.sqrt %20762 : tensor<1xf32> loc(#loc5249)
    %20764 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20765 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20766 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20767 = call @nan_to_num_2723(%20763, %20764, %20765, %20766) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20768 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20769 = stablehlo.broadcast_in_dim %20768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20770 = stablehlo.divide %20767, %20769 : tensor<1xf32> loc(#loc5253)
    %20771 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20772 = stablehlo.broadcast_in_dim %20771, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20773 = stablehlo.maximum %20772, %20770 : tensor<1xf32> loc(#loc5271)
    %20774 = stablehlo.broadcast_in_dim %20773, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %20775 = stablehlo.broadcast_in_dim %20774, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20776 = stablehlo.divide %17808, %20775 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20777 = stablehlo.multiply %17816, %17816 : tensor<1x64xf32> loc(#loc5268)
    %20778 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20779 = stablehlo.reduce(%20777 init: %20778) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20780 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %20781 = stablehlo.broadcast_in_dim %20780, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20782 = stablehlo.divide %20779, %20781 : tensor<1xf32> loc(#loc5270)
    %20783 = stablehlo.sqrt %20782 : tensor<1xf32> loc(#loc5249)
    %20784 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20785 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20786 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20787 = call @nan_to_num_2728(%20783, %20784, %20785, %20786) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20788 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20789 = stablehlo.broadcast_in_dim %20788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20790 = stablehlo.divide %20787, %20789 : tensor<1xf32> loc(#loc5253)
    %20791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20792 = stablehlo.broadcast_in_dim %20791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20793 = stablehlo.maximum %20792, %20790 : tensor<1xf32> loc(#loc5271)
    %20794 = stablehlo.broadcast_in_dim %20793, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20795 = stablehlo.broadcast_in_dim %20794, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %20796 = stablehlo.divide %17816, %20795 : tensor<1x64xf32> loc(#loc5272)
    %20797 = stablehlo.multiply %17824, %17824 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %20798 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %20799 = stablehlo.reduce(%20797 init: %20798) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %20800 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %20801 = stablehlo.broadcast_in_dim %20800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20802 = stablehlo.divide %20799, %20801 : tensor<1xf32> loc(#loc5270)
    %20803 = stablehlo.sqrt %20802 : tensor<1xf32> loc(#loc5249)
    %20804 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20805 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20806 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20807 = call @nan_to_num_2733(%20803, %20804, %20805, %20806) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20808 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20809 = stablehlo.broadcast_in_dim %20808, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20810 = stablehlo.divide %20807, %20809 : tensor<1xf32> loc(#loc5253)
    %20811 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20812 = stablehlo.broadcast_in_dim %20811, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20813 = stablehlo.maximum %20812, %20810 : tensor<1xf32> loc(#loc5271)
    %20814 = stablehlo.broadcast_in_dim %20813, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %20815 = stablehlo.broadcast_in_dim %20814, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %20816 = stablehlo.divide %17824, %20815 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %20817 = stablehlo.multiply %17832, %17832 : tensor<1x8192xf32> loc(#loc5268)
    %20818 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20819 = stablehlo.reduce(%20817 init: %20818) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20820 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %20821 = stablehlo.broadcast_in_dim %20820, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20822 = stablehlo.divide %20819, %20821 : tensor<1xf32> loc(#loc5270)
    %20823 = stablehlo.sqrt %20822 : tensor<1xf32> loc(#loc5249)
    %20824 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20825 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20826 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20827 = call @nan_to_num_2738(%20823, %20824, %20825, %20826) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20828 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20829 = stablehlo.broadcast_in_dim %20828, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20830 = stablehlo.divide %20827, %20829 : tensor<1xf32> loc(#loc5253)
    %20831 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20832 = stablehlo.broadcast_in_dim %20831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20833 = stablehlo.maximum %20832, %20830 : tensor<1xf32> loc(#loc5271)
    %20834 = stablehlo.broadcast_in_dim %20833, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20835 = stablehlo.broadcast_in_dim %20834, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %20836 = stablehlo.divide %17832, %20835 : tensor<1x8192xf32> loc(#loc5272)
    %20837 = stablehlo.multiply %17840, %17840 : tensor<1x2048x8192xf32> loc(#loc5268)
    %20838 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20839 = stablehlo.reduce(%20837 init: %20838) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20840 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20841 = stablehlo.broadcast_in_dim %20840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20842 = stablehlo.divide %20839, %20841 : tensor<1xf32> loc(#loc5270)
    %20843 = stablehlo.sqrt %20842 : tensor<1xf32> loc(#loc5249)
    %20844 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20845 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20846 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20847 = call @nan_to_num_2743(%20843, %20844, %20845, %20846) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20848 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20849 = stablehlo.broadcast_in_dim %20848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20850 = stablehlo.divide %20847, %20849 : tensor<1xf32> loc(#loc5253)
    %20851 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20852 = stablehlo.broadcast_in_dim %20851, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20853 = stablehlo.maximum %20852, %20850 : tensor<1xf32> loc(#loc5271)
    %20854 = stablehlo.broadcast_in_dim %20853, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20855 = stablehlo.broadcast_in_dim %20854, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %20856 = stablehlo.divide %17840, %20855 : tensor<1x2048x8192xf32> loc(#loc5272)
    %20857 = stablehlo.multiply %17848, %17848 : tensor<1x2048xf32> loc(#loc5268)
    %20858 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20859 = stablehlo.reduce(%20857 init: %20858) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20860 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20861 = stablehlo.broadcast_in_dim %20860, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20862 = stablehlo.divide %20859, %20861 : tensor<1xf32> loc(#loc5270)
    %20863 = stablehlo.sqrt %20862 : tensor<1xf32> loc(#loc5249)
    %20864 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20865 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20866 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20867 = call @nan_to_num_2748(%20863, %20864, %20865, %20866) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20868 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20869 = stablehlo.broadcast_in_dim %20868, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20870 = stablehlo.divide %20867, %20869 : tensor<1xf32> loc(#loc5253)
    %20871 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20872 = stablehlo.broadcast_in_dim %20871, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20873 = stablehlo.maximum %20872, %20870 : tensor<1xf32> loc(#loc5271)
    %20874 = stablehlo.broadcast_in_dim %20873, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20875 = stablehlo.broadcast_in_dim %20874, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20876 = stablehlo.divide %17848, %20875 : tensor<1x2048xf32> loc(#loc5272)
    %20877 = stablehlo.multiply %17856, %17856 : tensor<1x8192x2048xf32> loc(#loc5268)
    %20878 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %20879 = stablehlo.reduce(%20877 init: %20878) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %20880 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %20881 = stablehlo.broadcast_in_dim %20880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20882 = stablehlo.divide %20879, %20881 : tensor<1xf32> loc(#loc5270)
    %20883 = stablehlo.sqrt %20882 : tensor<1xf32> loc(#loc5249)
    %20884 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20885 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20886 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20887 = call @nan_to_num_2753(%20883, %20884, %20885, %20886) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20888 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20889 = stablehlo.broadcast_in_dim %20888, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20890 = stablehlo.divide %20887, %20889 : tensor<1xf32> loc(#loc5253)
    %20891 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20892 = stablehlo.broadcast_in_dim %20891, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20893 = stablehlo.maximum %20892, %20890 : tensor<1xf32> loc(#loc5271)
    %20894 = stablehlo.broadcast_in_dim %20893, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %20895 = stablehlo.broadcast_in_dim %20894, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %20896 = stablehlo.divide %17856, %20895 : tensor<1x8192x2048xf32> loc(#loc5272)
    %20897 = stablehlo.multiply %17864, %17864 : tensor<1x2048xf32> loc(#loc5268)
    %20898 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20899 = stablehlo.reduce(%20897 init: %20898) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20900 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20901 = stablehlo.broadcast_in_dim %20900, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20902 = stablehlo.divide %20899, %20901 : tensor<1xf32> loc(#loc5270)
    %20903 = stablehlo.sqrt %20902 : tensor<1xf32> loc(#loc5249)
    %20904 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20905 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20906 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20907 = call @nan_to_num_2758(%20903, %20904, %20905, %20906) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20908 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20909 = stablehlo.broadcast_in_dim %20908, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20910 = stablehlo.divide %20907, %20909 : tensor<1xf32> loc(#loc5253)
    %20911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20912 = stablehlo.broadcast_in_dim %20911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20913 = stablehlo.maximum %20912, %20910 : tensor<1xf32> loc(#loc5271)
    %20914 = stablehlo.broadcast_in_dim %20913, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20915 = stablehlo.broadcast_in_dim %20914, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20916 = stablehlo.divide %17864, %20915 : tensor<1x2048xf32> loc(#loc5272)
    %20917 = stablehlo.multiply %17872, %17872 : tensor<1x2048xf32> loc(#loc5268)
    %20918 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20919 = stablehlo.reduce(%20917 init: %20918) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20920 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20921 = stablehlo.broadcast_in_dim %20920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20922 = stablehlo.divide %20919, %20921 : tensor<1xf32> loc(#loc5270)
    %20923 = stablehlo.sqrt %20922 : tensor<1xf32> loc(#loc5249)
    %20924 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20925 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20926 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20927 = call @nan_to_num_2763(%20923, %20924, %20925, %20926) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20928 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20929 = stablehlo.broadcast_in_dim %20928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20930 = stablehlo.divide %20927, %20929 : tensor<1xf32> loc(#loc5253)
    %20931 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20932 = stablehlo.broadcast_in_dim %20931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20933 = stablehlo.maximum %20932, %20930 : tensor<1xf32> loc(#loc5271)
    %20934 = stablehlo.broadcast_in_dim %20933, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20935 = stablehlo.broadcast_in_dim %20934, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20936 = stablehlo.divide %17872, %20935 : tensor<1x2048xf32> loc(#loc5272)
    %20937 = stablehlo.multiply %17880, %17880 : tensor<1x2048xf32> loc(#loc5268)
    %20938 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20939 = stablehlo.reduce(%20937 init: %20938) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20940 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20941 = stablehlo.broadcast_in_dim %20940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20942 = stablehlo.divide %20939, %20941 : tensor<1xf32> loc(#loc5270)
    %20943 = stablehlo.sqrt %20942 : tensor<1xf32> loc(#loc5249)
    %20944 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20945 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20946 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20947 = call @nan_to_num_2768(%20943, %20944, %20945, %20946) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20948 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20949 = stablehlo.broadcast_in_dim %20948, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20950 = stablehlo.divide %20947, %20949 : tensor<1xf32> loc(#loc5253)
    %20951 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20952 = stablehlo.broadcast_in_dim %20951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20953 = stablehlo.maximum %20952, %20950 : tensor<1xf32> loc(#loc5271)
    %20954 = stablehlo.broadcast_in_dim %20953, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20955 = stablehlo.broadcast_in_dim %20954, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20956 = stablehlo.divide %17880, %20955 : tensor<1x2048xf32> loc(#loc5272)
    %20957 = stablehlo.multiply %17888, %17888 : tensor<1x2048xf32> loc(#loc5268)
    %20958 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20959 = stablehlo.reduce(%20957 init: %20958) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %20960 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %20961 = stablehlo.broadcast_in_dim %20960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20962 = stablehlo.divide %20959, %20961 : tensor<1xf32> loc(#loc5270)
    %20963 = stablehlo.sqrt %20962 : tensor<1xf32> loc(#loc5249)
    %20964 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20965 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20966 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20967 = call @nan_to_num_2773(%20963, %20964, %20965, %20966) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20968 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20969 = stablehlo.broadcast_in_dim %20968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20970 = stablehlo.divide %20967, %20969 : tensor<1xf32> loc(#loc5253)
    %20971 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20972 = stablehlo.broadcast_in_dim %20971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20973 = stablehlo.maximum %20972, %20970 : tensor<1xf32> loc(#loc5271)
    %20974 = stablehlo.broadcast_in_dim %20973, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %20975 = stablehlo.broadcast_in_dim %20974, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %20976 = stablehlo.divide %17888, %20975 : tensor<1x2048xf32> loc(#loc5272)
    %20977 = stablehlo.multiply %17896, %17896 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %20978 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %20979 = stablehlo.reduce(%20977 init: %20978) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %20980 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %20981 = stablehlo.broadcast_in_dim %20980, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %20982 = stablehlo.divide %20979, %20981 : tensor<1xf32> loc(#loc5270)
    %20983 = stablehlo.sqrt %20982 : tensor<1xf32> loc(#loc5249)
    %20984 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20985 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20986 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %20987 = call @nan_to_num_2778(%20983, %20984, %20985, %20986) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %20988 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20989 = stablehlo.broadcast_in_dim %20988, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %20990 = stablehlo.divide %20987, %20989 : tensor<1xf32> loc(#loc5253)
    %20991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %20992 = stablehlo.broadcast_in_dim %20991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %20993 = stablehlo.maximum %20992, %20990 : tensor<1xf32> loc(#loc5271)
    %20994 = stablehlo.broadcast_in_dim %20993, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %20995 = stablehlo.broadcast_in_dim %20994, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20996 = stablehlo.divide %17896, %20995 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %20997 = stablehlo.multiply %17904, %17904 : tensor<1x64xf32> loc(#loc5268)
    %20998 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %20999 = stablehlo.reduce(%20997 init: %20998) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21000 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %21001 = stablehlo.broadcast_in_dim %21000, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21002 = stablehlo.divide %20999, %21001 : tensor<1xf32> loc(#loc5270)
    %21003 = stablehlo.sqrt %21002 : tensor<1xf32> loc(#loc5249)
    %21004 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21005 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21006 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21007 = call @nan_to_num_2783(%21003, %21004, %21005, %21006) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21008 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21009 = stablehlo.broadcast_in_dim %21008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21010 = stablehlo.divide %21007, %21009 : tensor<1xf32> loc(#loc5253)
    %21011 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21012 = stablehlo.broadcast_in_dim %21011, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21013 = stablehlo.maximum %21012, %21010 : tensor<1xf32> loc(#loc5271)
    %21014 = stablehlo.broadcast_in_dim %21013, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21015 = stablehlo.broadcast_in_dim %21014, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %21016 = stablehlo.divide %17904, %21015 : tensor<1x64xf32> loc(#loc5272)
    %21017 = stablehlo.multiply %17912, %17912 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %21018 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %21019 = stablehlo.reduce(%21017 init: %21018) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %21020 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %21021 = stablehlo.broadcast_in_dim %21020, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21022 = stablehlo.divide %21019, %21021 : tensor<1xf32> loc(#loc5270)
    %21023 = stablehlo.sqrt %21022 : tensor<1xf32> loc(#loc5249)
    %21024 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21025 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21026 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21027 = call @nan_to_num_2788(%21023, %21024, %21025, %21026) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21028 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21029 = stablehlo.broadcast_in_dim %21028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21030 = stablehlo.divide %21027, %21029 : tensor<1xf32> loc(#loc5253)
    %21031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21032 = stablehlo.broadcast_in_dim %21031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21033 = stablehlo.maximum %21032, %21030 : tensor<1xf32> loc(#loc5271)
    %21034 = stablehlo.broadcast_in_dim %21033, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %21035 = stablehlo.broadcast_in_dim %21034, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %21036 = stablehlo.divide %17912, %21035 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %21037 = stablehlo.multiply %17920, %17920 : tensor<1x8192xf32> loc(#loc5268)
    %21038 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21039 = stablehlo.reduce(%21037 init: %21038) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21040 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %21041 = stablehlo.broadcast_in_dim %21040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21042 = stablehlo.divide %21039, %21041 : tensor<1xf32> loc(#loc5270)
    %21043 = stablehlo.sqrt %21042 : tensor<1xf32> loc(#loc5249)
    %21044 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21045 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21046 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21047 = call @nan_to_num_2793(%21043, %21044, %21045, %21046) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21048 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21049 = stablehlo.broadcast_in_dim %21048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21050 = stablehlo.divide %21047, %21049 : tensor<1xf32> loc(#loc5253)
    %21051 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21052 = stablehlo.broadcast_in_dim %21051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21053 = stablehlo.maximum %21052, %21050 : tensor<1xf32> loc(#loc5271)
    %21054 = stablehlo.broadcast_in_dim %21053, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21055 = stablehlo.broadcast_in_dim %21054, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %21056 = stablehlo.divide %17920, %21055 : tensor<1x8192xf32> loc(#loc5272)
    %21057 = stablehlo.multiply %17928, %17928 : tensor<1x2048x8192xf32> loc(#loc5268)
    %21058 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %21059 = stablehlo.reduce(%21057 init: %21058) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %21060 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %21061 = stablehlo.broadcast_in_dim %21060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21062 = stablehlo.divide %21059, %21061 : tensor<1xf32> loc(#loc5270)
    %21063 = stablehlo.sqrt %21062 : tensor<1xf32> loc(#loc5249)
    %21064 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21065 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21066 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21067 = call @nan_to_num_2798(%21063, %21064, %21065, %21066) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21068 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21069 = stablehlo.broadcast_in_dim %21068, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21070 = stablehlo.divide %21067, %21069 : tensor<1xf32> loc(#loc5253)
    %21071 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21072 = stablehlo.broadcast_in_dim %21071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21073 = stablehlo.maximum %21072, %21070 : tensor<1xf32> loc(#loc5271)
    %21074 = stablehlo.broadcast_in_dim %21073, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %21075 = stablehlo.broadcast_in_dim %21074, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %21076 = stablehlo.divide %17928, %21075 : tensor<1x2048x8192xf32> loc(#loc5272)
    %21077 = stablehlo.multiply %17936, %17936 : tensor<1x2048xf32> loc(#loc5268)
    %21078 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21079 = stablehlo.reduce(%21077 init: %21078) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21080 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21081 = stablehlo.broadcast_in_dim %21080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21082 = stablehlo.divide %21079, %21081 : tensor<1xf32> loc(#loc5270)
    %21083 = stablehlo.sqrt %21082 : tensor<1xf32> loc(#loc5249)
    %21084 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21085 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21086 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21087 = call @nan_to_num_2803(%21083, %21084, %21085, %21086) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21088 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21089 = stablehlo.broadcast_in_dim %21088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21090 = stablehlo.divide %21087, %21089 : tensor<1xf32> loc(#loc5253)
    %21091 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21092 = stablehlo.broadcast_in_dim %21091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21093 = stablehlo.maximum %21092, %21090 : tensor<1xf32> loc(#loc5271)
    %21094 = stablehlo.broadcast_in_dim %21093, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21095 = stablehlo.broadcast_in_dim %21094, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21096 = stablehlo.divide %17936, %21095 : tensor<1x2048xf32> loc(#loc5272)
    %21097 = stablehlo.multiply %17944, %17944 : tensor<1x8192x2048xf32> loc(#loc5268)
    %21098 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %21099 = stablehlo.reduce(%21097 init: %21098) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %21100 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %21101 = stablehlo.broadcast_in_dim %21100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21102 = stablehlo.divide %21099, %21101 : tensor<1xf32> loc(#loc5270)
    %21103 = stablehlo.sqrt %21102 : tensor<1xf32> loc(#loc5249)
    %21104 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21105 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21106 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21107 = call @nan_to_num_2808(%21103, %21104, %21105, %21106) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21108 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21109 = stablehlo.broadcast_in_dim %21108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21110 = stablehlo.divide %21107, %21109 : tensor<1xf32> loc(#loc5253)
    %21111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21112 = stablehlo.broadcast_in_dim %21111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21113 = stablehlo.maximum %21112, %21110 : tensor<1xf32> loc(#loc5271)
    %21114 = stablehlo.broadcast_in_dim %21113, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %21115 = stablehlo.broadcast_in_dim %21114, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %21116 = stablehlo.divide %17944, %21115 : tensor<1x8192x2048xf32> loc(#loc5272)
    %21117 = stablehlo.multiply %17952, %17952 : tensor<1x2048xf32> loc(#loc5268)
    %21118 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21119 = stablehlo.reduce(%21117 init: %21118) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21120 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21121 = stablehlo.broadcast_in_dim %21120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21122 = stablehlo.divide %21119, %21121 : tensor<1xf32> loc(#loc5270)
    %21123 = stablehlo.sqrt %21122 : tensor<1xf32> loc(#loc5249)
    %21124 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21125 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21126 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21127 = call @nan_to_num_2813(%21123, %21124, %21125, %21126) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21128 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21129 = stablehlo.broadcast_in_dim %21128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21130 = stablehlo.divide %21127, %21129 : tensor<1xf32> loc(#loc5253)
    %21131 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21132 = stablehlo.broadcast_in_dim %21131, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21133 = stablehlo.maximum %21132, %21130 : tensor<1xf32> loc(#loc5271)
    %21134 = stablehlo.broadcast_in_dim %21133, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21135 = stablehlo.broadcast_in_dim %21134, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21136 = stablehlo.divide %17952, %21135 : tensor<1x2048xf32> loc(#loc5272)
    %21137 = stablehlo.multiply %17960, %17960 : tensor<1x2048xf32> loc(#loc5268)
    %21138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21139 = stablehlo.reduce(%21137 init: %21138) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21140 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21141 = stablehlo.broadcast_in_dim %21140, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21142 = stablehlo.divide %21139, %21141 : tensor<1xf32> loc(#loc5270)
    %21143 = stablehlo.sqrt %21142 : tensor<1xf32> loc(#loc5249)
    %21144 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21145 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21146 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21147 = call @nan_to_num_2818(%21143, %21144, %21145, %21146) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21148 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21149 = stablehlo.broadcast_in_dim %21148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21150 = stablehlo.divide %21147, %21149 : tensor<1xf32> loc(#loc5253)
    %21151 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21152 = stablehlo.broadcast_in_dim %21151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21153 = stablehlo.maximum %21152, %21150 : tensor<1xf32> loc(#loc5271)
    %21154 = stablehlo.broadcast_in_dim %21153, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21155 = stablehlo.broadcast_in_dim %21154, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21156 = stablehlo.divide %17960, %21155 : tensor<1x2048xf32> loc(#loc5272)
    %21157 = stablehlo.multiply %17968, %17968 : tensor<1x2048xf32> loc(#loc5268)
    %21158 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21159 = stablehlo.reduce(%21157 init: %21158) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21160 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21161 = stablehlo.broadcast_in_dim %21160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21162 = stablehlo.divide %21159, %21161 : tensor<1xf32> loc(#loc5270)
    %21163 = stablehlo.sqrt %21162 : tensor<1xf32> loc(#loc5249)
    %21164 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21165 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21166 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21167 = call @nan_to_num_2823(%21163, %21164, %21165, %21166) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21168 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21169 = stablehlo.broadcast_in_dim %21168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21170 = stablehlo.divide %21167, %21169 : tensor<1xf32> loc(#loc5253)
    %21171 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21172 = stablehlo.broadcast_in_dim %21171, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21173 = stablehlo.maximum %21172, %21170 : tensor<1xf32> loc(#loc5271)
    %21174 = stablehlo.broadcast_in_dim %21173, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21175 = stablehlo.broadcast_in_dim %21174, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21176 = stablehlo.divide %17968, %21175 : tensor<1x2048xf32> loc(#loc5272)
    %21177 = stablehlo.multiply %17976, %17976 : tensor<1x2048xf32> loc(#loc5268)
    %21178 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21179 = stablehlo.reduce(%21177 init: %21178) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21180 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21181 = stablehlo.broadcast_in_dim %21180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21182 = stablehlo.divide %21179, %21181 : tensor<1xf32> loc(#loc5270)
    %21183 = stablehlo.sqrt %21182 : tensor<1xf32> loc(#loc5249)
    %21184 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21185 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21186 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21187 = call @nan_to_num_2828(%21183, %21184, %21185, %21186) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21188 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21189 = stablehlo.broadcast_in_dim %21188, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21190 = stablehlo.divide %21187, %21189 : tensor<1xf32> loc(#loc5253)
    %21191 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21192 = stablehlo.broadcast_in_dim %21191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21193 = stablehlo.maximum %21192, %21190 : tensor<1xf32> loc(#loc5271)
    %21194 = stablehlo.broadcast_in_dim %21193, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21195 = stablehlo.broadcast_in_dim %21194, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21196 = stablehlo.divide %17976, %21195 : tensor<1x2048xf32> loc(#loc5272)
    %21197 = stablehlo.multiply %17984, %17984 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %21198 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %21199 = stablehlo.reduce(%21197 init: %21198) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %21200 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %21201 = stablehlo.broadcast_in_dim %21200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21202 = stablehlo.divide %21199, %21201 : tensor<1xf32> loc(#loc5270)
    %21203 = stablehlo.sqrt %21202 : tensor<1xf32> loc(#loc5249)
    %21204 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21205 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21206 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21207 = call @nan_to_num_2833(%21203, %21204, %21205, %21206) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21208 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21209 = stablehlo.broadcast_in_dim %21208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21210 = stablehlo.divide %21207, %21209 : tensor<1xf32> loc(#loc5253)
    %21211 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21212 = stablehlo.broadcast_in_dim %21211, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21213 = stablehlo.maximum %21212, %21210 : tensor<1xf32> loc(#loc5271)
    %21214 = stablehlo.broadcast_in_dim %21213, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %21215 = stablehlo.broadcast_in_dim %21214, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %21216 = stablehlo.divide %17984, %21215 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %21217 = stablehlo.multiply %17992, %17992 : tensor<1x64xf32> loc(#loc5268)
    %21218 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21219 = stablehlo.reduce(%21217 init: %21218) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21220 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %21221 = stablehlo.broadcast_in_dim %21220, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21222 = stablehlo.divide %21219, %21221 : tensor<1xf32> loc(#loc5270)
    %21223 = stablehlo.sqrt %21222 : tensor<1xf32> loc(#loc5249)
    %21224 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21225 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21226 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21227 = call @nan_to_num_2838(%21223, %21224, %21225, %21226) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21228 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21229 = stablehlo.broadcast_in_dim %21228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21230 = stablehlo.divide %21227, %21229 : tensor<1xf32> loc(#loc5253)
    %21231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21232 = stablehlo.broadcast_in_dim %21231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21233 = stablehlo.maximum %21232, %21230 : tensor<1xf32> loc(#loc5271)
    %21234 = stablehlo.broadcast_in_dim %21233, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21235 = stablehlo.broadcast_in_dim %21234, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %21236 = stablehlo.divide %17992, %21235 : tensor<1x64xf32> loc(#loc5272)
    %21237 = stablehlo.multiply %18000, %18000 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %21238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %21239 = stablehlo.reduce(%21237 init: %21238) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %21240 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %21241 = stablehlo.broadcast_in_dim %21240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21242 = stablehlo.divide %21239, %21241 : tensor<1xf32> loc(#loc5270)
    %21243 = stablehlo.sqrt %21242 : tensor<1xf32> loc(#loc5249)
    %21244 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21245 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21246 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21247 = call @nan_to_num_2843(%21243, %21244, %21245, %21246) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21248 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21249 = stablehlo.broadcast_in_dim %21248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21250 = stablehlo.divide %21247, %21249 : tensor<1xf32> loc(#loc5253)
    %21251 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21252 = stablehlo.broadcast_in_dim %21251, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21253 = stablehlo.maximum %21252, %21250 : tensor<1xf32> loc(#loc5271)
    %21254 = stablehlo.broadcast_in_dim %21253, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %21255 = stablehlo.broadcast_in_dim %21254, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %21256 = stablehlo.divide %18000, %21255 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %21257 = stablehlo.multiply %18008, %18008 : tensor<1x8192xf32> loc(#loc5268)
    %21258 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21259 = stablehlo.reduce(%21257 init: %21258) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21260 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %21261 = stablehlo.broadcast_in_dim %21260, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21262 = stablehlo.divide %21259, %21261 : tensor<1xf32> loc(#loc5270)
    %21263 = stablehlo.sqrt %21262 : tensor<1xf32> loc(#loc5249)
    %21264 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21265 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21266 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21267 = call @nan_to_num_2848(%21263, %21264, %21265, %21266) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21268 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21269 = stablehlo.broadcast_in_dim %21268, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21270 = stablehlo.divide %21267, %21269 : tensor<1xf32> loc(#loc5253)
    %21271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21272 = stablehlo.broadcast_in_dim %21271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21273 = stablehlo.maximum %21272, %21270 : tensor<1xf32> loc(#loc5271)
    %21274 = stablehlo.broadcast_in_dim %21273, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21275 = stablehlo.broadcast_in_dim %21274, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %21276 = stablehlo.divide %18008, %21275 : tensor<1x8192xf32> loc(#loc5272)
    %21277 = stablehlo.multiply %18016, %18016 : tensor<1x2048x8192xf32> loc(#loc5268)
    %21278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %21279 = stablehlo.reduce(%21277 init: %21278) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %21280 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %21281 = stablehlo.broadcast_in_dim %21280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21282 = stablehlo.divide %21279, %21281 : tensor<1xf32> loc(#loc5270)
    %21283 = stablehlo.sqrt %21282 : tensor<1xf32> loc(#loc5249)
    %21284 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21285 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21286 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21287 = call @nan_to_num_2853(%21283, %21284, %21285, %21286) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21288 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21289 = stablehlo.broadcast_in_dim %21288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21290 = stablehlo.divide %21287, %21289 : tensor<1xf32> loc(#loc5253)
    %21291 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21292 = stablehlo.broadcast_in_dim %21291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21293 = stablehlo.maximum %21292, %21290 : tensor<1xf32> loc(#loc5271)
    %21294 = stablehlo.broadcast_in_dim %21293, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %21295 = stablehlo.broadcast_in_dim %21294, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %21296 = stablehlo.divide %18016, %21295 : tensor<1x2048x8192xf32> loc(#loc5272)
    %21297 = stablehlo.multiply %18024, %18024 : tensor<1x2048xf32> loc(#loc5268)
    %21298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21299 = stablehlo.reduce(%21297 init: %21298) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21300 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21301 = stablehlo.broadcast_in_dim %21300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21302 = stablehlo.divide %21299, %21301 : tensor<1xf32> loc(#loc5270)
    %21303 = stablehlo.sqrt %21302 : tensor<1xf32> loc(#loc5249)
    %21304 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21305 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21306 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21307 = call @nan_to_num_2858(%21303, %21304, %21305, %21306) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21308 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21309 = stablehlo.broadcast_in_dim %21308, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21310 = stablehlo.divide %21307, %21309 : tensor<1xf32> loc(#loc5253)
    %21311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21312 = stablehlo.broadcast_in_dim %21311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21313 = stablehlo.maximum %21312, %21310 : tensor<1xf32> loc(#loc5271)
    %21314 = stablehlo.broadcast_in_dim %21313, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21315 = stablehlo.broadcast_in_dim %21314, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21316 = stablehlo.divide %18024, %21315 : tensor<1x2048xf32> loc(#loc5272)
    %21317 = stablehlo.multiply %18032, %18032 : tensor<1x8192x2048xf32> loc(#loc5268)
    %21318 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %21319 = stablehlo.reduce(%21317 init: %21318) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %21320 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %21321 = stablehlo.broadcast_in_dim %21320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21322 = stablehlo.divide %21319, %21321 : tensor<1xf32> loc(#loc5270)
    %21323 = stablehlo.sqrt %21322 : tensor<1xf32> loc(#loc5249)
    %21324 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21325 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21326 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21327 = call @nan_to_num_2863(%21323, %21324, %21325, %21326) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21328 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21329 = stablehlo.broadcast_in_dim %21328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21330 = stablehlo.divide %21327, %21329 : tensor<1xf32> loc(#loc5253)
    %21331 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21332 = stablehlo.broadcast_in_dim %21331, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21333 = stablehlo.maximum %21332, %21330 : tensor<1xf32> loc(#loc5271)
    %21334 = stablehlo.broadcast_in_dim %21333, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %21335 = stablehlo.broadcast_in_dim %21334, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %21336 = stablehlo.divide %18032, %21335 : tensor<1x8192x2048xf32> loc(#loc5272)
    %21337 = stablehlo.multiply %18040, %18040 : tensor<1x2048xf32> loc(#loc5268)
    %21338 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21339 = stablehlo.reduce(%21337 init: %21338) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21340 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21341 = stablehlo.broadcast_in_dim %21340, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21342 = stablehlo.divide %21339, %21341 : tensor<1xf32> loc(#loc5270)
    %21343 = stablehlo.sqrt %21342 : tensor<1xf32> loc(#loc5249)
    %21344 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21345 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21346 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21347 = call @nan_to_num_2868(%21343, %21344, %21345, %21346) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21348 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21349 = stablehlo.broadcast_in_dim %21348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21350 = stablehlo.divide %21347, %21349 : tensor<1xf32> loc(#loc5253)
    %21351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21352 = stablehlo.broadcast_in_dim %21351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21353 = stablehlo.maximum %21352, %21350 : tensor<1xf32> loc(#loc5271)
    %21354 = stablehlo.broadcast_in_dim %21353, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21355 = stablehlo.broadcast_in_dim %21354, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21356 = stablehlo.divide %18040, %21355 : tensor<1x2048xf32> loc(#loc5272)
    %21357 = stablehlo.multiply %18048, %18048 : tensor<1x2048xf32> loc(#loc5268)
    %21358 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21359 = stablehlo.reduce(%21357 init: %21358) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21360 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21361 = stablehlo.broadcast_in_dim %21360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21362 = stablehlo.divide %21359, %21361 : tensor<1xf32> loc(#loc5270)
    %21363 = stablehlo.sqrt %21362 : tensor<1xf32> loc(#loc5249)
    %21364 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21365 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21366 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21367 = call @nan_to_num_2873(%21363, %21364, %21365, %21366) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21368 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21369 = stablehlo.broadcast_in_dim %21368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21370 = stablehlo.divide %21367, %21369 : tensor<1xf32> loc(#loc5253)
    %21371 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21372 = stablehlo.broadcast_in_dim %21371, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21373 = stablehlo.maximum %21372, %21370 : tensor<1xf32> loc(#loc5271)
    %21374 = stablehlo.broadcast_in_dim %21373, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21375 = stablehlo.broadcast_in_dim %21374, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21376 = stablehlo.divide %18048, %21375 : tensor<1x2048xf32> loc(#loc5272)
    %21377 = stablehlo.multiply %18056, %18056 : tensor<1x2048xf32> loc(#loc5268)
    %21378 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21379 = stablehlo.reduce(%21377 init: %21378) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21380 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21381 = stablehlo.broadcast_in_dim %21380, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21382 = stablehlo.divide %21379, %21381 : tensor<1xf32> loc(#loc5270)
    %21383 = stablehlo.sqrt %21382 : tensor<1xf32> loc(#loc5249)
    %21384 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21385 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21386 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21387 = call @nan_to_num_2878(%21383, %21384, %21385, %21386) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21388 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21389 = stablehlo.broadcast_in_dim %21388, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21390 = stablehlo.divide %21387, %21389 : tensor<1xf32> loc(#loc5253)
    %21391 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21392 = stablehlo.broadcast_in_dim %21391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21393 = stablehlo.maximum %21392, %21390 : tensor<1xf32> loc(#loc5271)
    %21394 = stablehlo.broadcast_in_dim %21393, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21395 = stablehlo.broadcast_in_dim %21394, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21396 = stablehlo.divide %18056, %21395 : tensor<1x2048xf32> loc(#loc5272)
    %21397 = stablehlo.multiply %18064, %18064 : tensor<1x2048xf32> loc(#loc5268)
    %21398 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21399 = stablehlo.reduce(%21397 init: %21398) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21400 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21401 = stablehlo.broadcast_in_dim %21400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21402 = stablehlo.divide %21399, %21401 : tensor<1xf32> loc(#loc5270)
    %21403 = stablehlo.sqrt %21402 : tensor<1xf32> loc(#loc5249)
    %21404 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21405 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21406 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21407 = call @nan_to_num_2883(%21403, %21404, %21405, %21406) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21408 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21409 = stablehlo.broadcast_in_dim %21408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21410 = stablehlo.divide %21407, %21409 : tensor<1xf32> loc(#loc5253)
    %21411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21412 = stablehlo.broadcast_in_dim %21411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21413 = stablehlo.maximum %21412, %21410 : tensor<1xf32> loc(#loc5271)
    %21414 = stablehlo.broadcast_in_dim %21413, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21415 = stablehlo.broadcast_in_dim %21414, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21416 = stablehlo.divide %18064, %21415 : tensor<1x2048xf32> loc(#loc5272)
    %21417 = stablehlo.multiply %18072, %18072 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %21418 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %21419 = stablehlo.reduce(%21417 init: %21418) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %21420 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %21421 = stablehlo.broadcast_in_dim %21420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21422 = stablehlo.divide %21419, %21421 : tensor<1xf32> loc(#loc5270)
    %21423 = stablehlo.sqrt %21422 : tensor<1xf32> loc(#loc5249)
    %21424 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21425 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21426 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21427 = call @nan_to_num_2888(%21423, %21424, %21425, %21426) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21428 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21429 = stablehlo.broadcast_in_dim %21428, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21430 = stablehlo.divide %21427, %21429 : tensor<1xf32> loc(#loc5253)
    %21431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21432 = stablehlo.broadcast_in_dim %21431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21433 = stablehlo.maximum %21432, %21430 : tensor<1xf32> loc(#loc5271)
    %21434 = stablehlo.broadcast_in_dim %21433, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %21435 = stablehlo.broadcast_in_dim %21434, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %21436 = stablehlo.divide %18072, %21435 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %21437 = stablehlo.multiply %18080, %18080 : tensor<1x64xf32> loc(#loc5268)
    %21438 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21439 = stablehlo.reduce(%21437 init: %21438) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21440 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %21441 = stablehlo.broadcast_in_dim %21440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21442 = stablehlo.divide %21439, %21441 : tensor<1xf32> loc(#loc5270)
    %21443 = stablehlo.sqrt %21442 : tensor<1xf32> loc(#loc5249)
    %21444 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21445 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21446 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21447 = call @nan_to_num_2893(%21443, %21444, %21445, %21446) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21448 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21449 = stablehlo.broadcast_in_dim %21448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21450 = stablehlo.divide %21447, %21449 : tensor<1xf32> loc(#loc5253)
    %21451 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21452 = stablehlo.broadcast_in_dim %21451, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21453 = stablehlo.maximum %21452, %21450 : tensor<1xf32> loc(#loc5271)
    %21454 = stablehlo.broadcast_in_dim %21453, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21455 = stablehlo.broadcast_in_dim %21454, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %21456 = stablehlo.divide %18080, %21455 : tensor<1x64xf32> loc(#loc5272)
    %21457 = stablehlo.multiply %18088, %18088 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %21458 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %21459 = stablehlo.reduce(%21457 init: %21458) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %21460 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %21461 = stablehlo.broadcast_in_dim %21460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21462 = stablehlo.divide %21459, %21461 : tensor<1xf32> loc(#loc5270)
    %21463 = stablehlo.sqrt %21462 : tensor<1xf32> loc(#loc5249)
    %21464 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21465 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21466 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21467 = call @nan_to_num_2898(%21463, %21464, %21465, %21466) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21468 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21469 = stablehlo.broadcast_in_dim %21468, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21470 = stablehlo.divide %21467, %21469 : tensor<1xf32> loc(#loc5253)
    %21471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21472 = stablehlo.broadcast_in_dim %21471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21473 = stablehlo.maximum %21472, %21470 : tensor<1xf32> loc(#loc5271)
    %21474 = stablehlo.broadcast_in_dim %21473, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %21475 = stablehlo.broadcast_in_dim %21474, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %21476 = stablehlo.divide %18088, %21475 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %21477 = stablehlo.multiply %18096, %18096 : tensor<1x8192xf32> loc(#loc5268)
    %21478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21479 = stablehlo.reduce(%21477 init: %21478) across dimensions = [1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21480 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %21481 = stablehlo.broadcast_in_dim %21480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21482 = stablehlo.divide %21479, %21481 : tensor<1xf32> loc(#loc5270)
    %21483 = stablehlo.sqrt %21482 : tensor<1xf32> loc(#loc5249)
    %21484 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21485 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21486 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21487 = call @nan_to_num_2903(%21483, %21484, %21485, %21486) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21488 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21489 = stablehlo.broadcast_in_dim %21488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21490 = stablehlo.divide %21487, %21489 : tensor<1xf32> loc(#loc5253)
    %21491 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21492 = stablehlo.broadcast_in_dim %21491, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21493 = stablehlo.maximum %21492, %21490 : tensor<1xf32> loc(#loc5271)
    %21494 = stablehlo.broadcast_in_dim %21493, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21495 = stablehlo.broadcast_in_dim %21494, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5272)
    %21496 = stablehlo.divide %18096, %21495 : tensor<1x8192xf32> loc(#loc5272)
    %21497 = stablehlo.multiply %18104, %18104 : tensor<1x2048x8192xf32> loc(#loc5268)
    %21498 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %21499 = stablehlo.reduce(%21497 init: %21498) across dimensions = [1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %21500 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %21501 = stablehlo.broadcast_in_dim %21500, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21502 = stablehlo.divide %21499, %21501 : tensor<1xf32> loc(#loc5270)
    %21503 = stablehlo.sqrt %21502 : tensor<1xf32> loc(#loc5249)
    %21504 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21505 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21506 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21507 = call @nan_to_num_2908(%21503, %21504, %21505, %21506) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21508 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21509 = stablehlo.broadcast_in_dim %21508, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21510 = stablehlo.divide %21507, %21509 : tensor<1xf32> loc(#loc5253)
    %21511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21512 = stablehlo.broadcast_in_dim %21511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21513 = stablehlo.maximum %21512, %21510 : tensor<1xf32> loc(#loc5271)
    %21514 = stablehlo.broadcast_in_dim %21513, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %21515 = stablehlo.broadcast_in_dim %21514, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5272)
    %21516 = stablehlo.divide %18104, %21515 : tensor<1x2048x8192xf32> loc(#loc5272)
    %21517 = stablehlo.multiply %18112, %18112 : tensor<1x2048xf32> loc(#loc5268)
    %21518 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21519 = stablehlo.reduce(%21517 init: %21518) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21520 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21521 = stablehlo.broadcast_in_dim %21520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21522 = stablehlo.divide %21519, %21521 : tensor<1xf32> loc(#loc5270)
    %21523 = stablehlo.sqrt %21522 : tensor<1xf32> loc(#loc5249)
    %21524 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21525 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21526 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21527 = call @nan_to_num_2913(%21523, %21524, %21525, %21526) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21528 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21529 = stablehlo.broadcast_in_dim %21528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21530 = stablehlo.divide %21527, %21529 : tensor<1xf32> loc(#loc5253)
    %21531 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21532 = stablehlo.broadcast_in_dim %21531, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21533 = stablehlo.maximum %21532, %21530 : tensor<1xf32> loc(#loc5271)
    %21534 = stablehlo.broadcast_in_dim %21533, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21535 = stablehlo.broadcast_in_dim %21534, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21536 = stablehlo.divide %18112, %21535 : tensor<1x2048xf32> loc(#loc5272)
    %21537 = stablehlo.multiply %18120, %18120 : tensor<1x8192x2048xf32> loc(#loc5268)
    %21538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5284)
    %21539 = stablehlo.reduce(%21537 init: %21538) across dimensions = [1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5284)
      stablehlo.return %29929 : tensor<f32> loc(#loc5284)
    } loc(#loc5284)
    %21540 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %21541 = stablehlo.broadcast_in_dim %21540, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21542 = stablehlo.divide %21539, %21541 : tensor<1xf32> loc(#loc5270)
    %21543 = stablehlo.sqrt %21542 : tensor<1xf32> loc(#loc5249)
    %21544 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21545 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21546 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21547 = call @nan_to_num_2918(%21543, %21544, %21545, %21546) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21548 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21549 = stablehlo.broadcast_in_dim %21548, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21550 = stablehlo.divide %21547, %21549 : tensor<1xf32> loc(#loc5253)
    %21551 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21552 = stablehlo.broadcast_in_dim %21551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21553 = stablehlo.maximum %21552, %21550 : tensor<1xf32> loc(#loc5271)
    %21554 = stablehlo.broadcast_in_dim %21553, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5285)
    %21555 = stablehlo.broadcast_in_dim %21554, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5272)
    %21556 = stablehlo.divide %18120, %21555 : tensor<1x8192x2048xf32> loc(#loc5272)
    %21557 = stablehlo.multiply %18128, %18128 : tensor<1x2048xf32> loc(#loc5268)
    %21558 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21559 = stablehlo.reduce(%21557 init: %21558) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21560 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21561 = stablehlo.broadcast_in_dim %21560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21562 = stablehlo.divide %21559, %21561 : tensor<1xf32> loc(#loc5270)
    %21563 = stablehlo.sqrt %21562 : tensor<1xf32> loc(#loc5249)
    %21564 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21565 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21566 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21567 = call @nan_to_num_2923(%21563, %21564, %21565, %21566) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21568 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21569 = stablehlo.broadcast_in_dim %21568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21570 = stablehlo.divide %21567, %21569 : tensor<1xf32> loc(#loc5253)
    %21571 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21572 = stablehlo.broadcast_in_dim %21571, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21573 = stablehlo.maximum %21572, %21570 : tensor<1xf32> loc(#loc5271)
    %21574 = stablehlo.broadcast_in_dim %21573, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21575 = stablehlo.broadcast_in_dim %21574, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21576 = stablehlo.divide %18128, %21575 : tensor<1x2048xf32> loc(#loc5272)
    %21577 = stablehlo.multiply %18136, %18136 : tensor<1x2048xf32> loc(#loc5268)
    %21578 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21579 = stablehlo.reduce(%21577 init: %21578) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21580 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21581 = stablehlo.broadcast_in_dim %21580, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21582 = stablehlo.divide %21579, %21581 : tensor<1xf32> loc(#loc5270)
    %21583 = stablehlo.sqrt %21582 : tensor<1xf32> loc(#loc5249)
    %21584 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21585 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21586 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21587 = call @nan_to_num_2928(%21583, %21584, %21585, %21586) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21588 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21589 = stablehlo.broadcast_in_dim %21588, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21590 = stablehlo.divide %21587, %21589 : tensor<1xf32> loc(#loc5253)
    %21591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21592 = stablehlo.broadcast_in_dim %21591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21593 = stablehlo.maximum %21592, %21590 : tensor<1xf32> loc(#loc5271)
    %21594 = stablehlo.broadcast_in_dim %21593, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21595 = stablehlo.broadcast_in_dim %21594, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21596 = stablehlo.divide %18136, %21595 : tensor<1x2048xf32> loc(#loc5272)
    %21597 = stablehlo.multiply %18144, %18144 : tensor<1x2048xf32> loc(#loc5268)
    %21598 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21599 = stablehlo.reduce(%21597 init: %21598) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21600 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21601 = stablehlo.broadcast_in_dim %21600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21602 = stablehlo.divide %21599, %21601 : tensor<1xf32> loc(#loc5270)
    %21603 = stablehlo.sqrt %21602 : tensor<1xf32> loc(#loc5249)
    %21604 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21605 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21606 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21607 = call @nan_to_num_2933(%21603, %21604, %21605, %21606) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21608 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21609 = stablehlo.broadcast_in_dim %21608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21610 = stablehlo.divide %21607, %21609 : tensor<1xf32> loc(#loc5253)
    %21611 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21612 = stablehlo.broadcast_in_dim %21611, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21613 = stablehlo.maximum %21612, %21610 : tensor<1xf32> loc(#loc5271)
    %21614 = stablehlo.broadcast_in_dim %21613, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21615 = stablehlo.broadcast_in_dim %21614, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21616 = stablehlo.divide %18144, %21615 : tensor<1x2048xf32> loc(#loc5272)
    %21617 = stablehlo.multiply %18152, %18152 : tensor<1x2048xf32> loc(#loc5268)
    %21618 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21619 = stablehlo.reduce(%21617 init: %21618) across dimensions = [1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21620 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %21621 = stablehlo.broadcast_in_dim %21620, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21622 = stablehlo.divide %21619, %21621 : tensor<1xf32> loc(#loc5270)
    %21623 = stablehlo.sqrt %21622 : tensor<1xf32> loc(#loc5249)
    %21624 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21625 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21626 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21627 = call @nan_to_num_2938(%21623, %21624, %21625, %21626) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21628 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21629 = stablehlo.broadcast_in_dim %21628, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21630 = stablehlo.divide %21627, %21629 : tensor<1xf32> loc(#loc5253)
    %21631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21632 = stablehlo.broadcast_in_dim %21631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21633 = stablehlo.maximum %21632, %21630 : tensor<1xf32> loc(#loc5271)
    %21634 = stablehlo.broadcast_in_dim %21633, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21635 = stablehlo.broadcast_in_dim %21634, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5272)
    %21636 = stablehlo.divide %18152, %21635 : tensor<1x2048xf32> loc(#loc5272)
    %21637 = stablehlo.multiply %18160, %18160 : tensor<1x3x2048x32x64xf32> loc(#loc5268)
    %21638 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5286)
    %21639 = stablehlo.reduce(%21637 init: %21638) across dimensions = [1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5286)
      stablehlo.return %29929 : tensor<f32> loc(#loc5286)
    } loc(#loc5286)
    %21640 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %21641 = stablehlo.broadcast_in_dim %21640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21642 = stablehlo.divide %21639, %21641 : tensor<1xf32> loc(#loc5270)
    %21643 = stablehlo.sqrt %21642 : tensor<1xf32> loc(#loc5249)
    %21644 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21645 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21646 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21647 = call @nan_to_num_2943(%21643, %21644, %21645, %21646) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21648 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21649 = stablehlo.broadcast_in_dim %21648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21650 = stablehlo.divide %21647, %21649 : tensor<1xf32> loc(#loc5253)
    %21651 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21652 = stablehlo.broadcast_in_dim %21651, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21653 = stablehlo.maximum %21652, %21650 : tensor<1xf32> loc(#loc5271)
    %21654 = stablehlo.broadcast_in_dim %21653, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5287)
    %21655 = stablehlo.broadcast_in_dim %21654, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %21656 = stablehlo.divide %18160, %21655 : tensor<1x3x2048x32x64xf32> loc(#loc5272)
    %21657 = stablehlo.multiply %18168, %18168 : tensor<1x64xf32> loc(#loc5268)
    %21658 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5282)
    %21659 = stablehlo.reduce(%21657 init: %21658) across dimensions = [1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5282)
      stablehlo.return %29929 : tensor<f32> loc(#loc5282)
    } loc(#loc5282)
    %21660 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %21661 = stablehlo.broadcast_in_dim %21660, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21662 = stablehlo.divide %21659, %21661 : tensor<1xf32> loc(#loc5270)
    %21663 = stablehlo.sqrt %21662 : tensor<1xf32> loc(#loc5249)
    %21664 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21665 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21666 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21667 = call @nan_to_num_2948(%21663, %21664, %21665, %21666) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21668 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21669 = stablehlo.broadcast_in_dim %21668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21670 = stablehlo.divide %21667, %21669 : tensor<1xf32> loc(#loc5253)
    %21671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21672 = stablehlo.broadcast_in_dim %21671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21673 = stablehlo.maximum %21672, %21670 : tensor<1xf32> loc(#loc5271)
    %21674 = stablehlo.broadcast_in_dim %21673, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5283)
    %21675 = stablehlo.broadcast_in_dim %21674, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5272)
    %21676 = stablehlo.divide %18168, %21675 : tensor<1x64xf32> loc(#loc5272)
    %21677 = stablehlo.multiply %18176, %18176 : tensor<1x2048x32x64xf32> loc(#loc5268)
    %21678 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5288)
    %21679 = stablehlo.reduce(%21677 init: %21678) across dimensions = [1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<1xf32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5288)
      stablehlo.return %29929 : tensor<f32> loc(#loc5288)
    } loc(#loc5288)
    %21680 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %21681 = stablehlo.broadcast_in_dim %21680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5270)
    %21682 = stablehlo.divide %21679, %21681 : tensor<1xf32> loc(#loc5270)
    %21683 = stablehlo.sqrt %21682 : tensor<1xf32> loc(#loc5249)
    %21684 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21685 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21686 = stablehlo.constant dense<0x7FC00000> : tensor<f32> loc(#loc)
    %21687 = call @nan_to_num_2953(%21683, %21684, %21685, %21686) : (tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<1xf32> loc(#loc5257)
    %21688 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21689 = stablehlo.broadcast_in_dim %21688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5253)
    %21690 = stablehlo.divide %21687, %21689 : tensor<1xf32> loc(#loc5253)
    %21691 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %21692 = stablehlo.broadcast_in_dim %21691, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5271)
    %21693 = stablehlo.maximum %21692, %21690 : tensor<1xf32> loc(#loc5271)
    %21694 = stablehlo.broadcast_in_dim %21693, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5289)
    %21695 = stablehlo.broadcast_in_dim %21694, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5272)
    %21696 = stablehlo.divide %18176, %21695 : tensor<1x2048x32x64xf32> loc(#loc5272)
    %21697 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21698 = stablehlo.broadcast_in_dim %21697, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21699 = stablehlo.multiply %21698, %arg6 : tensor<1x8192xf32> loc(#loc5275)
    %21700 = stablehlo.add %18196, %21699 : tensor<1x8192xf32> loc(#loc5263)
    %21701 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21702 = stablehlo.broadcast_in_dim %21701, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21703 = stablehlo.multiply %21702, %arg7 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21704 = stablehlo.add %18216, %21703 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21705 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21706 = stablehlo.broadcast_in_dim %21705, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21707 = stablehlo.multiply %21706, %arg8 : tensor<1x2048xf32> loc(#loc5275)
    %21708 = stablehlo.add %18236, %21707 : tensor<1x2048xf32> loc(#loc5263)
    %21709 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21710 = stablehlo.broadcast_in_dim %21709, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21711 = stablehlo.multiply %21710, %arg9 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21712 = stablehlo.add %18256, %21711 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21713 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21714 = stablehlo.broadcast_in_dim %21713, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21715 = stablehlo.multiply %21714, %arg10 : tensor<1x2048xf32> loc(#loc5275)
    %21716 = stablehlo.add %18276, %21715 : tensor<1x2048xf32> loc(#loc5263)
    %21717 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21718 = stablehlo.broadcast_in_dim %21717, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21719 = stablehlo.multiply %21718, %arg11 : tensor<1x2048xf32> loc(#loc5275)
    %21720 = stablehlo.add %18296, %21719 : tensor<1x2048xf32> loc(#loc5263)
    %21721 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21722 = stablehlo.broadcast_in_dim %21721, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21723 = stablehlo.multiply %21722, %arg12 : tensor<1x2048xf32> loc(#loc5275)
    %21724 = stablehlo.add %18316, %21723 : tensor<1x2048xf32> loc(#loc5263)
    %21725 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21726 = stablehlo.broadcast_in_dim %21725, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21727 = stablehlo.multiply %21726, %arg13 : tensor<1x2048xf32> loc(#loc5275)
    %21728 = stablehlo.add %18336, %21727 : tensor<1x2048xf32> loc(#loc5263)
    %21729 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21730 = stablehlo.broadcast_in_dim %21729, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21731 = stablehlo.multiply %21730, %arg14 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21732 = stablehlo.add %18356, %21731 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21733 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21734 = stablehlo.broadcast_in_dim %21733, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21735 = stablehlo.multiply %21734, %arg15 : tensor<1x64xf32> loc(#loc5275)
    %21736 = stablehlo.add %18376, %21735 : tensor<1x64xf32> loc(#loc5263)
    %21737 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21738 = stablehlo.broadcast_in_dim %21737, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %21739 = stablehlo.multiply %21738, %arg16 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %21740 = stablehlo.add %18396, %21739 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %21741 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21742 = stablehlo.broadcast_in_dim %21741, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21743 = stablehlo.multiply %21742, %arg17 : tensor<1x8192xf32> loc(#loc5275)
    %21744 = stablehlo.add %18416, %21743 : tensor<1x8192xf32> loc(#loc5263)
    %21745 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21746 = stablehlo.broadcast_in_dim %21745, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21747 = stablehlo.multiply %21746, %arg18 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21748 = stablehlo.add %18436, %21747 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21749 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21750 = stablehlo.broadcast_in_dim %21749, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21751 = stablehlo.multiply %21750, %arg19 : tensor<1x2048xf32> loc(#loc5275)
    %21752 = stablehlo.add %18456, %21751 : tensor<1x2048xf32> loc(#loc5263)
    %21753 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21754 = stablehlo.broadcast_in_dim %21753, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21755 = stablehlo.multiply %21754, %arg20 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21756 = stablehlo.add %18476, %21755 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21757 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21758 = stablehlo.broadcast_in_dim %21757, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21759 = stablehlo.multiply %21758, %arg21 : tensor<1x2048xf32> loc(#loc5275)
    %21760 = stablehlo.add %18496, %21759 : tensor<1x2048xf32> loc(#loc5263)
    %21761 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21762 = stablehlo.broadcast_in_dim %21761, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21763 = stablehlo.multiply %21762, %arg22 : tensor<1x2048xf32> loc(#loc5275)
    %21764 = stablehlo.add %18516, %21763 : tensor<1x2048xf32> loc(#loc5263)
    %21765 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21766 = stablehlo.broadcast_in_dim %21765, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21767 = stablehlo.multiply %21766, %arg23 : tensor<1x2048xf32> loc(#loc5275)
    %21768 = stablehlo.add %18536, %21767 : tensor<1x2048xf32> loc(#loc5263)
    %21769 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21770 = stablehlo.broadcast_in_dim %21769, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21771 = stablehlo.multiply %21770, %arg24 : tensor<1x2048xf32> loc(#loc5275)
    %21772 = stablehlo.add %18556, %21771 : tensor<1x2048xf32> loc(#loc5263)
    %21773 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21774 = stablehlo.broadcast_in_dim %21773, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21775 = stablehlo.multiply %21774, %arg25 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21776 = stablehlo.add %18576, %21775 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21777 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21778 = stablehlo.broadcast_in_dim %21777, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21779 = stablehlo.multiply %21778, %arg26 : tensor<1x64xf32> loc(#loc5275)
    %21780 = stablehlo.add %18596, %21779 : tensor<1x64xf32> loc(#loc5263)
    %21781 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21782 = stablehlo.broadcast_in_dim %21781, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %21783 = stablehlo.multiply %21782, %arg27 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %21784 = stablehlo.add %18616, %21783 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %21785 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21786 = stablehlo.broadcast_in_dim %21785, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21787 = stablehlo.multiply %21786, %arg28 : tensor<1x8192xf32> loc(#loc5275)
    %21788 = stablehlo.add %18636, %21787 : tensor<1x8192xf32> loc(#loc5263)
    %21789 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21790 = stablehlo.broadcast_in_dim %21789, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21791 = stablehlo.multiply %21790, %arg29 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21792 = stablehlo.add %18656, %21791 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21793 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21794 = stablehlo.broadcast_in_dim %21793, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21795 = stablehlo.multiply %21794, %arg30 : tensor<1x2048xf32> loc(#loc5275)
    %21796 = stablehlo.add %18676, %21795 : tensor<1x2048xf32> loc(#loc5263)
    %21797 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21798 = stablehlo.broadcast_in_dim %21797, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21799 = stablehlo.multiply %21798, %arg31 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21800 = stablehlo.add %18696, %21799 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21801 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21802 = stablehlo.broadcast_in_dim %21801, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21803 = stablehlo.multiply %21802, %arg32 : tensor<1x2048xf32> loc(#loc5275)
    %21804 = stablehlo.add %18716, %21803 : tensor<1x2048xf32> loc(#loc5263)
    %21805 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21806 = stablehlo.broadcast_in_dim %21805, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21807 = stablehlo.multiply %21806, %arg33 : tensor<1x2048xf32> loc(#loc5275)
    %21808 = stablehlo.add %18736, %21807 : tensor<1x2048xf32> loc(#loc5263)
    %21809 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21810 = stablehlo.broadcast_in_dim %21809, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21811 = stablehlo.multiply %21810, %arg34 : tensor<1x2048xf32> loc(#loc5275)
    %21812 = stablehlo.add %18756, %21811 : tensor<1x2048xf32> loc(#loc5263)
    %21813 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21814 = stablehlo.broadcast_in_dim %21813, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21815 = stablehlo.multiply %21814, %arg35 : tensor<1x2048xf32> loc(#loc5275)
    %21816 = stablehlo.add %18776, %21815 : tensor<1x2048xf32> loc(#loc5263)
    %21817 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21818 = stablehlo.broadcast_in_dim %21817, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21819 = stablehlo.multiply %21818, %arg36 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21820 = stablehlo.add %18796, %21819 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21821 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21822 = stablehlo.broadcast_in_dim %21821, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21823 = stablehlo.multiply %21822, %arg37 : tensor<1x64xf32> loc(#loc5275)
    %21824 = stablehlo.add %18816, %21823 : tensor<1x64xf32> loc(#loc5263)
    %21825 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21826 = stablehlo.broadcast_in_dim %21825, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %21827 = stablehlo.multiply %21826, %arg38 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %21828 = stablehlo.add %18836, %21827 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %21829 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21830 = stablehlo.broadcast_in_dim %21829, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21831 = stablehlo.multiply %21830, %arg39 : tensor<1x8192xf32> loc(#loc5275)
    %21832 = stablehlo.add %18856, %21831 : tensor<1x8192xf32> loc(#loc5263)
    %21833 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21834 = stablehlo.broadcast_in_dim %21833, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21835 = stablehlo.multiply %21834, %arg40 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21836 = stablehlo.add %18876, %21835 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21837 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21838 = stablehlo.broadcast_in_dim %21837, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21839 = stablehlo.multiply %21838, %arg41 : tensor<1x2048xf32> loc(#loc5275)
    %21840 = stablehlo.add %18896, %21839 : tensor<1x2048xf32> loc(#loc5263)
    %21841 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21842 = stablehlo.broadcast_in_dim %21841, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21843 = stablehlo.multiply %21842, %arg42 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21844 = stablehlo.add %18916, %21843 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21845 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21846 = stablehlo.broadcast_in_dim %21845, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21847 = stablehlo.multiply %21846, %arg43 : tensor<1x2048xf32> loc(#loc5275)
    %21848 = stablehlo.add %18936, %21847 : tensor<1x2048xf32> loc(#loc5263)
    %21849 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21850 = stablehlo.broadcast_in_dim %21849, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21851 = stablehlo.multiply %21850, %arg44 : tensor<1x2048xf32> loc(#loc5275)
    %21852 = stablehlo.add %18956, %21851 : tensor<1x2048xf32> loc(#loc5263)
    %21853 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21854 = stablehlo.broadcast_in_dim %21853, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21855 = stablehlo.multiply %21854, %arg45 : tensor<1x2048xf32> loc(#loc5275)
    %21856 = stablehlo.add %18976, %21855 : tensor<1x2048xf32> loc(#loc5263)
    %21857 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21858 = stablehlo.broadcast_in_dim %21857, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21859 = stablehlo.multiply %21858, %arg46 : tensor<1x2048xf32> loc(#loc5275)
    %21860 = stablehlo.add %18996, %21859 : tensor<1x2048xf32> loc(#loc5263)
    %21861 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21862 = stablehlo.broadcast_in_dim %21861, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21863 = stablehlo.multiply %21862, %arg47 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21864 = stablehlo.add %19016, %21863 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21865 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21866 = stablehlo.broadcast_in_dim %21865, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21867 = stablehlo.multiply %21866, %arg48 : tensor<1x64xf32> loc(#loc5275)
    %21868 = stablehlo.add %19036, %21867 : tensor<1x64xf32> loc(#loc5263)
    %21869 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21870 = stablehlo.broadcast_in_dim %21869, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %21871 = stablehlo.multiply %21870, %arg49 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %21872 = stablehlo.add %19056, %21871 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %21873 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21874 = stablehlo.broadcast_in_dim %21873, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21875 = stablehlo.multiply %21874, %arg50 : tensor<1x8192xf32> loc(#loc5275)
    %21876 = stablehlo.add %19076, %21875 : tensor<1x8192xf32> loc(#loc5263)
    %21877 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21878 = stablehlo.broadcast_in_dim %21877, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21879 = stablehlo.multiply %21878, %arg51 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21880 = stablehlo.add %19096, %21879 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21881 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21882 = stablehlo.broadcast_in_dim %21881, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21883 = stablehlo.multiply %21882, %arg52 : tensor<1x2048xf32> loc(#loc5275)
    %21884 = stablehlo.add %19116, %21883 : tensor<1x2048xf32> loc(#loc5263)
    %21885 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21886 = stablehlo.broadcast_in_dim %21885, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21887 = stablehlo.multiply %21886, %arg53 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21888 = stablehlo.add %19136, %21887 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21889 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21890 = stablehlo.broadcast_in_dim %21889, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21891 = stablehlo.multiply %21890, %arg54 : tensor<1x2048xf32> loc(#loc5275)
    %21892 = stablehlo.add %19156, %21891 : tensor<1x2048xf32> loc(#loc5263)
    %21893 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21894 = stablehlo.broadcast_in_dim %21893, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21895 = stablehlo.multiply %21894, %arg55 : tensor<1x2048xf32> loc(#loc5275)
    %21896 = stablehlo.add %19176, %21895 : tensor<1x2048xf32> loc(#loc5263)
    %21897 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21898 = stablehlo.broadcast_in_dim %21897, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21899 = stablehlo.multiply %21898, %arg56 : tensor<1x2048xf32> loc(#loc5275)
    %21900 = stablehlo.add %19196, %21899 : tensor<1x2048xf32> loc(#loc5263)
    %21901 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21902 = stablehlo.broadcast_in_dim %21901, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21903 = stablehlo.multiply %21902, %arg57 : tensor<1x2048xf32> loc(#loc5275)
    %21904 = stablehlo.add %19216, %21903 : tensor<1x2048xf32> loc(#loc5263)
    %21905 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21906 = stablehlo.broadcast_in_dim %21905, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21907 = stablehlo.multiply %21906, %arg58 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21908 = stablehlo.add %19236, %21907 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21909 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21910 = stablehlo.broadcast_in_dim %21909, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21911 = stablehlo.multiply %21910, %arg59 : tensor<1x64xf32> loc(#loc5275)
    %21912 = stablehlo.add %19256, %21911 : tensor<1x64xf32> loc(#loc5263)
    %21913 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21914 = stablehlo.broadcast_in_dim %21913, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %21915 = stablehlo.multiply %21914, %arg60 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %21916 = stablehlo.add %19276, %21915 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %21917 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21918 = stablehlo.broadcast_in_dim %21917, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21919 = stablehlo.multiply %21918, %arg61 : tensor<1x8192xf32> loc(#loc5275)
    %21920 = stablehlo.add %19296, %21919 : tensor<1x8192xf32> loc(#loc5263)
    %21921 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21922 = stablehlo.broadcast_in_dim %21921, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21923 = stablehlo.multiply %21922, %arg62 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21924 = stablehlo.add %19316, %21923 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21925 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21926 = stablehlo.broadcast_in_dim %21925, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21927 = stablehlo.multiply %21926, %arg63 : tensor<1x2048xf32> loc(#loc5275)
    %21928 = stablehlo.add %19336, %21927 : tensor<1x2048xf32> loc(#loc5263)
    %21929 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21930 = stablehlo.broadcast_in_dim %21929, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21931 = stablehlo.multiply %21930, %arg64 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21932 = stablehlo.add %19356, %21931 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21933 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21934 = stablehlo.broadcast_in_dim %21933, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21935 = stablehlo.multiply %21934, %arg65 : tensor<1x2048xf32> loc(#loc5275)
    %21936 = stablehlo.add %19376, %21935 : tensor<1x2048xf32> loc(#loc5263)
    %21937 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21938 = stablehlo.broadcast_in_dim %21937, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21939 = stablehlo.multiply %21938, %arg66 : tensor<1x2048xf32> loc(#loc5275)
    %21940 = stablehlo.add %19396, %21939 : tensor<1x2048xf32> loc(#loc5263)
    %21941 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21942 = stablehlo.broadcast_in_dim %21941, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21943 = stablehlo.multiply %21942, %arg67 : tensor<1x2048xf32> loc(#loc5275)
    %21944 = stablehlo.add %19416, %21943 : tensor<1x2048xf32> loc(#loc5263)
    %21945 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21946 = stablehlo.broadcast_in_dim %21945, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21947 = stablehlo.multiply %21946, %arg68 : tensor<1x2048xf32> loc(#loc5275)
    %21948 = stablehlo.add %19436, %21947 : tensor<1x2048xf32> loc(#loc5263)
    %21949 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21950 = stablehlo.broadcast_in_dim %21949, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21951 = stablehlo.multiply %21950, %arg69 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21952 = stablehlo.add %19456, %21951 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21953 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21954 = stablehlo.broadcast_in_dim %21953, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21955 = stablehlo.multiply %21954, %arg70 : tensor<1x64xf32> loc(#loc5275)
    %21956 = stablehlo.add %19476, %21955 : tensor<1x64xf32> loc(#loc5263)
    %21957 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21958 = stablehlo.broadcast_in_dim %21957, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %21959 = stablehlo.multiply %21958, %arg71 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %21960 = stablehlo.add %19496, %21959 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %21961 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21962 = stablehlo.broadcast_in_dim %21961, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %21963 = stablehlo.multiply %21962, %arg72 : tensor<1x8192xf32> loc(#loc5275)
    %21964 = stablehlo.add %19516, %21963 : tensor<1x8192xf32> loc(#loc5263)
    %21965 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21966 = stablehlo.broadcast_in_dim %21965, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %21967 = stablehlo.multiply %21966, %arg73 : tensor<1x2048x8192xf32> loc(#loc5276)
    %21968 = stablehlo.add %19536, %21967 : tensor<1x2048x8192xf32> loc(#loc5263)
    %21969 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21970 = stablehlo.broadcast_in_dim %21969, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21971 = stablehlo.multiply %21970, %arg74 : tensor<1x2048xf32> loc(#loc5275)
    %21972 = stablehlo.add %19556, %21971 : tensor<1x2048xf32> loc(#loc5263)
    %21973 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21974 = stablehlo.broadcast_in_dim %21973, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %21975 = stablehlo.multiply %21974, %arg75 : tensor<1x8192x2048xf32> loc(#loc5276)
    %21976 = stablehlo.add %19576, %21975 : tensor<1x8192x2048xf32> loc(#loc5263)
    %21977 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21978 = stablehlo.broadcast_in_dim %21977, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21979 = stablehlo.multiply %21978, %arg76 : tensor<1x2048xf32> loc(#loc5275)
    %21980 = stablehlo.add %19596, %21979 : tensor<1x2048xf32> loc(#loc5263)
    %21981 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21982 = stablehlo.broadcast_in_dim %21981, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21983 = stablehlo.multiply %21982, %arg77 : tensor<1x2048xf32> loc(#loc5275)
    %21984 = stablehlo.add %19616, %21983 : tensor<1x2048xf32> loc(#loc5263)
    %21985 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21986 = stablehlo.broadcast_in_dim %21985, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21987 = stablehlo.multiply %21986, %arg78 : tensor<1x2048xf32> loc(#loc5275)
    %21988 = stablehlo.add %19636, %21987 : tensor<1x2048xf32> loc(#loc5263)
    %21989 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21990 = stablehlo.broadcast_in_dim %21989, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %21991 = stablehlo.multiply %21990, %arg79 : tensor<1x2048xf32> loc(#loc5275)
    %21992 = stablehlo.add %19656, %21991 : tensor<1x2048xf32> loc(#loc5263)
    %21993 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21994 = stablehlo.broadcast_in_dim %21993, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21995 = stablehlo.multiply %21994, %arg80 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %21996 = stablehlo.add %19676, %21995 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %21997 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %21998 = stablehlo.broadcast_in_dim %21997, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %21999 = stablehlo.multiply %21998, %arg81 : tensor<1x64xf32> loc(#loc5275)
    %22000 = stablehlo.add %19696, %21999 : tensor<1x64xf32> loc(#loc5263)
    %22001 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22002 = stablehlo.broadcast_in_dim %22001, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22003 = stablehlo.multiply %22002, %arg82 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22004 = stablehlo.add %19716, %22003 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22005 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22006 = stablehlo.broadcast_in_dim %22005, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22007 = stablehlo.multiply %22006, %arg83 : tensor<1x8192xf32> loc(#loc5275)
    %22008 = stablehlo.add %19736, %22007 : tensor<1x8192xf32> loc(#loc5263)
    %22009 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22010 = stablehlo.broadcast_in_dim %22009, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22011 = stablehlo.multiply %22010, %arg84 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22012 = stablehlo.add %19756, %22011 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22013 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22014 = stablehlo.broadcast_in_dim %22013, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22015 = stablehlo.multiply %22014, %arg85 : tensor<1x2048xf32> loc(#loc5275)
    %22016 = stablehlo.add %19776, %22015 : tensor<1x2048xf32> loc(#loc5263)
    %22017 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22018 = stablehlo.broadcast_in_dim %22017, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22019 = stablehlo.multiply %22018, %arg86 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22020 = stablehlo.add %19796, %22019 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22021 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22022 = stablehlo.broadcast_in_dim %22021, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22023 = stablehlo.multiply %22022, %arg87 : tensor<1x2048xf32> loc(#loc5275)
    %22024 = stablehlo.add %19816, %22023 : tensor<1x2048xf32> loc(#loc5263)
    %22025 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22026 = stablehlo.broadcast_in_dim %22025, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22027 = stablehlo.multiply %22026, %arg88 : tensor<1x2048xf32> loc(#loc5275)
    %22028 = stablehlo.add %19836, %22027 : tensor<1x2048xf32> loc(#loc5263)
    %22029 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22030 = stablehlo.broadcast_in_dim %22029, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22031 = stablehlo.multiply %22030, %arg89 : tensor<1x2048xf32> loc(#loc5275)
    %22032 = stablehlo.add %19856, %22031 : tensor<1x2048xf32> loc(#loc5263)
    %22033 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22034 = stablehlo.broadcast_in_dim %22033, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22035 = stablehlo.multiply %22034, %arg90 : tensor<1x2048xf32> loc(#loc5275)
    %22036 = stablehlo.add %19876, %22035 : tensor<1x2048xf32> loc(#loc5263)
    %22037 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22038 = stablehlo.broadcast_in_dim %22037, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22039 = stablehlo.multiply %22038, %arg91 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22040 = stablehlo.add %19896, %22039 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22041 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22042 = stablehlo.broadcast_in_dim %22041, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22043 = stablehlo.multiply %22042, %arg92 : tensor<1x64xf32> loc(#loc5275)
    %22044 = stablehlo.add %19916, %22043 : tensor<1x64xf32> loc(#loc5263)
    %22045 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22046 = stablehlo.broadcast_in_dim %22045, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22047 = stablehlo.multiply %22046, %arg93 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22048 = stablehlo.add %19936, %22047 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22049 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22050 = stablehlo.broadcast_in_dim %22049, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22051 = stablehlo.multiply %22050, %arg94 : tensor<1x8192xf32> loc(#loc5275)
    %22052 = stablehlo.add %19956, %22051 : tensor<1x8192xf32> loc(#loc5263)
    %22053 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22054 = stablehlo.broadcast_in_dim %22053, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22055 = stablehlo.multiply %22054, %arg95 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22056 = stablehlo.add %19976, %22055 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22057 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22058 = stablehlo.broadcast_in_dim %22057, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22059 = stablehlo.multiply %22058, %arg96 : tensor<1x2048xf32> loc(#loc5275)
    %22060 = stablehlo.add %19996, %22059 : tensor<1x2048xf32> loc(#loc5263)
    %22061 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22062 = stablehlo.broadcast_in_dim %22061, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22063 = stablehlo.multiply %22062, %arg97 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22064 = stablehlo.add %20016, %22063 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22065 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22066 = stablehlo.broadcast_in_dim %22065, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22067 = stablehlo.multiply %22066, %arg98 : tensor<1x2048xf32> loc(#loc5275)
    %22068 = stablehlo.add %20036, %22067 : tensor<1x2048xf32> loc(#loc5263)
    %22069 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22070 = stablehlo.broadcast_in_dim %22069, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22071 = stablehlo.multiply %22070, %arg99 : tensor<1x2048xf32> loc(#loc5275)
    %22072 = stablehlo.add %20056, %22071 : tensor<1x2048xf32> loc(#loc5263)
    %22073 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22074 = stablehlo.broadcast_in_dim %22073, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22075 = stablehlo.multiply %22074, %arg100 : tensor<1x2048xf32> loc(#loc5275)
    %22076 = stablehlo.add %20076, %22075 : tensor<1x2048xf32> loc(#loc5263)
    %22077 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22078 = stablehlo.broadcast_in_dim %22077, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22079 = stablehlo.multiply %22078, %arg101 : tensor<1x2048xf32> loc(#loc5275)
    %22080 = stablehlo.add %20096, %22079 : tensor<1x2048xf32> loc(#loc5263)
    %22081 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22082 = stablehlo.broadcast_in_dim %22081, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22083 = stablehlo.multiply %22082, %arg102 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22084 = stablehlo.add %20116, %22083 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22085 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22086 = stablehlo.broadcast_in_dim %22085, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22087 = stablehlo.multiply %22086, %arg103 : tensor<1x64xf32> loc(#loc5275)
    %22088 = stablehlo.add %20136, %22087 : tensor<1x64xf32> loc(#loc5263)
    %22089 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22090 = stablehlo.broadcast_in_dim %22089, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22091 = stablehlo.multiply %22090, %arg104 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22092 = stablehlo.add %20156, %22091 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22093 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22094 = stablehlo.broadcast_in_dim %22093, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22095 = stablehlo.multiply %22094, %arg105 : tensor<1x8192xf32> loc(#loc5275)
    %22096 = stablehlo.add %20176, %22095 : tensor<1x8192xf32> loc(#loc5263)
    %22097 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22098 = stablehlo.broadcast_in_dim %22097, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22099 = stablehlo.multiply %22098, %arg106 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22100 = stablehlo.add %20196, %22099 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22101 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22102 = stablehlo.broadcast_in_dim %22101, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22103 = stablehlo.multiply %22102, %arg107 : tensor<1x2048xf32> loc(#loc5275)
    %22104 = stablehlo.add %20216, %22103 : tensor<1x2048xf32> loc(#loc5263)
    %22105 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22106 = stablehlo.broadcast_in_dim %22105, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22107 = stablehlo.multiply %22106, %arg108 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22108 = stablehlo.add %20236, %22107 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22109 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22110 = stablehlo.broadcast_in_dim %22109, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22111 = stablehlo.multiply %22110, %arg109 : tensor<1x2048xf32> loc(#loc5275)
    %22112 = stablehlo.add %20256, %22111 : tensor<1x2048xf32> loc(#loc5263)
    %22113 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22114 = stablehlo.broadcast_in_dim %22113, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22115 = stablehlo.multiply %22114, %arg110 : tensor<1x2048xf32> loc(#loc5275)
    %22116 = stablehlo.add %20276, %22115 : tensor<1x2048xf32> loc(#loc5263)
    %22117 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22118 = stablehlo.broadcast_in_dim %22117, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22119 = stablehlo.multiply %22118, %arg111 : tensor<1x2048xf32> loc(#loc5275)
    %22120 = stablehlo.add %20296, %22119 : tensor<1x2048xf32> loc(#loc5263)
    %22121 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22122 = stablehlo.broadcast_in_dim %22121, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22123 = stablehlo.multiply %22122, %arg112 : tensor<1x2048xf32> loc(#loc5275)
    %22124 = stablehlo.add %20316, %22123 : tensor<1x2048xf32> loc(#loc5263)
    %22125 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22126 = stablehlo.broadcast_in_dim %22125, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22127 = stablehlo.multiply %22126, %arg113 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22128 = stablehlo.add %20336, %22127 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22129 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22130 = stablehlo.broadcast_in_dim %22129, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22131 = stablehlo.multiply %22130, %arg114 : tensor<1x64xf32> loc(#loc5275)
    %22132 = stablehlo.add %20356, %22131 : tensor<1x64xf32> loc(#loc5263)
    %22133 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22134 = stablehlo.broadcast_in_dim %22133, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22135 = stablehlo.multiply %22134, %arg115 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22136 = stablehlo.add %20376, %22135 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22137 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22138 = stablehlo.broadcast_in_dim %22137, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22139 = stablehlo.multiply %22138, %arg116 : tensor<1x8192xf32> loc(#loc5275)
    %22140 = stablehlo.add %20396, %22139 : tensor<1x8192xf32> loc(#loc5263)
    %22141 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22142 = stablehlo.broadcast_in_dim %22141, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22143 = stablehlo.multiply %22142, %arg117 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22144 = stablehlo.add %20416, %22143 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22145 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22146 = stablehlo.broadcast_in_dim %22145, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22147 = stablehlo.multiply %22146, %arg118 : tensor<1x2048xf32> loc(#loc5275)
    %22148 = stablehlo.add %20436, %22147 : tensor<1x2048xf32> loc(#loc5263)
    %22149 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22150 = stablehlo.broadcast_in_dim %22149, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22151 = stablehlo.multiply %22150, %arg119 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22152 = stablehlo.add %20456, %22151 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22153 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22154 = stablehlo.broadcast_in_dim %22153, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22155 = stablehlo.multiply %22154, %arg120 : tensor<1x2048xf32> loc(#loc5275)
    %22156 = stablehlo.add %20476, %22155 : tensor<1x2048xf32> loc(#loc5263)
    %22157 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22158 = stablehlo.broadcast_in_dim %22157, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22159 = stablehlo.multiply %22158, %arg121 : tensor<1x2048xf32> loc(#loc5275)
    %22160 = stablehlo.add %20496, %22159 : tensor<1x2048xf32> loc(#loc5263)
    %22161 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22162 = stablehlo.broadcast_in_dim %22161, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22163 = stablehlo.multiply %22162, %arg122 : tensor<1x2048xf32> loc(#loc5275)
    %22164 = stablehlo.add %20516, %22163 : tensor<1x2048xf32> loc(#loc5263)
    %22165 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22166 = stablehlo.broadcast_in_dim %22165, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22167 = stablehlo.multiply %22166, %arg123 : tensor<1x2048xf32> loc(#loc5275)
    %22168 = stablehlo.add %20536, %22167 : tensor<1x2048xf32> loc(#loc5263)
    %22169 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22170 = stablehlo.broadcast_in_dim %22169, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22171 = stablehlo.multiply %22170, %arg124 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22172 = stablehlo.add %20556, %22171 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22173 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22174 = stablehlo.broadcast_in_dim %22173, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22175 = stablehlo.multiply %22174, %arg125 : tensor<1x64xf32> loc(#loc5275)
    %22176 = stablehlo.add %20576, %22175 : tensor<1x64xf32> loc(#loc5263)
    %22177 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22178 = stablehlo.broadcast_in_dim %22177, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22179 = stablehlo.multiply %22178, %arg126 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22180 = stablehlo.add %20596, %22179 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22181 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22182 = stablehlo.broadcast_in_dim %22181, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22183 = stablehlo.multiply %22182, %arg127 : tensor<1x8192xf32> loc(#loc5275)
    %22184 = stablehlo.add %20616, %22183 : tensor<1x8192xf32> loc(#loc5263)
    %22185 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22186 = stablehlo.broadcast_in_dim %22185, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22187 = stablehlo.multiply %22186, %arg128 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22188 = stablehlo.add %20636, %22187 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22189 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22190 = stablehlo.broadcast_in_dim %22189, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22191 = stablehlo.multiply %22190, %arg129 : tensor<1x2048xf32> loc(#loc5275)
    %22192 = stablehlo.add %20656, %22191 : tensor<1x2048xf32> loc(#loc5263)
    %22193 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22194 = stablehlo.broadcast_in_dim %22193, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22195 = stablehlo.multiply %22194, %arg130 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22196 = stablehlo.add %20676, %22195 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22197 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22198 = stablehlo.broadcast_in_dim %22197, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22199 = stablehlo.multiply %22198, %arg131 : tensor<1x2048xf32> loc(#loc5275)
    %22200 = stablehlo.add %20696, %22199 : tensor<1x2048xf32> loc(#loc5263)
    %22201 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22202 = stablehlo.broadcast_in_dim %22201, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22203 = stablehlo.multiply %22202, %arg132 : tensor<1x2048xf32> loc(#loc5275)
    %22204 = stablehlo.add %20716, %22203 : tensor<1x2048xf32> loc(#loc5263)
    %22205 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22206 = stablehlo.broadcast_in_dim %22205, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22207 = stablehlo.multiply %22206, %arg133 : tensor<1x2048xf32> loc(#loc5275)
    %22208 = stablehlo.add %20736, %22207 : tensor<1x2048xf32> loc(#loc5263)
    %22209 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22210 = stablehlo.broadcast_in_dim %22209, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22211 = stablehlo.multiply %22210, %arg134 : tensor<1x2048xf32> loc(#loc5275)
    %22212 = stablehlo.add %20756, %22211 : tensor<1x2048xf32> loc(#loc5263)
    %22213 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22214 = stablehlo.broadcast_in_dim %22213, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22215 = stablehlo.multiply %22214, %arg135 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22216 = stablehlo.add %20776, %22215 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22217 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22218 = stablehlo.broadcast_in_dim %22217, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22219 = stablehlo.multiply %22218, %arg136 : tensor<1x64xf32> loc(#loc5275)
    %22220 = stablehlo.add %20796, %22219 : tensor<1x64xf32> loc(#loc5263)
    %22221 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22222 = stablehlo.broadcast_in_dim %22221, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22223 = stablehlo.multiply %22222, %arg137 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22224 = stablehlo.add %20816, %22223 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22225 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22226 = stablehlo.broadcast_in_dim %22225, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22227 = stablehlo.multiply %22226, %arg138 : tensor<1x8192xf32> loc(#loc5275)
    %22228 = stablehlo.add %20836, %22227 : tensor<1x8192xf32> loc(#loc5263)
    %22229 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22230 = stablehlo.broadcast_in_dim %22229, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22231 = stablehlo.multiply %22230, %arg139 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22232 = stablehlo.add %20856, %22231 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22233 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22234 = stablehlo.broadcast_in_dim %22233, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22235 = stablehlo.multiply %22234, %arg140 : tensor<1x2048xf32> loc(#loc5275)
    %22236 = stablehlo.add %20876, %22235 : tensor<1x2048xf32> loc(#loc5263)
    %22237 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22238 = stablehlo.broadcast_in_dim %22237, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22239 = stablehlo.multiply %22238, %arg141 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22240 = stablehlo.add %20896, %22239 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22241 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22242 = stablehlo.broadcast_in_dim %22241, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22243 = stablehlo.multiply %22242, %arg142 : tensor<1x2048xf32> loc(#loc5275)
    %22244 = stablehlo.add %20916, %22243 : tensor<1x2048xf32> loc(#loc5263)
    %22245 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22246 = stablehlo.broadcast_in_dim %22245, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22247 = stablehlo.multiply %22246, %arg143 : tensor<1x2048xf32> loc(#loc5275)
    %22248 = stablehlo.add %20936, %22247 : tensor<1x2048xf32> loc(#loc5263)
    %22249 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22250 = stablehlo.broadcast_in_dim %22249, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22251 = stablehlo.multiply %22250, %arg144 : tensor<1x2048xf32> loc(#loc5275)
    %22252 = stablehlo.add %20956, %22251 : tensor<1x2048xf32> loc(#loc5263)
    %22253 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22254 = stablehlo.broadcast_in_dim %22253, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22255 = stablehlo.multiply %22254, %arg145 : tensor<1x2048xf32> loc(#loc5275)
    %22256 = stablehlo.add %20976, %22255 : tensor<1x2048xf32> loc(#loc5263)
    %22257 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22258 = stablehlo.broadcast_in_dim %22257, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22259 = stablehlo.multiply %22258, %arg146 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22260 = stablehlo.add %20996, %22259 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22261 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22262 = stablehlo.broadcast_in_dim %22261, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22263 = stablehlo.multiply %22262, %arg147 : tensor<1x64xf32> loc(#loc5275)
    %22264 = stablehlo.add %21016, %22263 : tensor<1x64xf32> loc(#loc5263)
    %22265 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22266 = stablehlo.broadcast_in_dim %22265, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22267 = stablehlo.multiply %22266, %arg148 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22268 = stablehlo.add %21036, %22267 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22269 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22270 = stablehlo.broadcast_in_dim %22269, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22271 = stablehlo.multiply %22270, %arg149 : tensor<1x8192xf32> loc(#loc5275)
    %22272 = stablehlo.add %21056, %22271 : tensor<1x8192xf32> loc(#loc5263)
    %22273 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22274 = stablehlo.broadcast_in_dim %22273, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22275 = stablehlo.multiply %22274, %arg150 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22276 = stablehlo.add %21076, %22275 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22277 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22278 = stablehlo.broadcast_in_dim %22277, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22279 = stablehlo.multiply %22278, %arg151 : tensor<1x2048xf32> loc(#loc5275)
    %22280 = stablehlo.add %21096, %22279 : tensor<1x2048xf32> loc(#loc5263)
    %22281 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22282 = stablehlo.broadcast_in_dim %22281, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22283 = stablehlo.multiply %22282, %arg152 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22284 = stablehlo.add %21116, %22283 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22285 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22286 = stablehlo.broadcast_in_dim %22285, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22287 = stablehlo.multiply %22286, %arg153 : tensor<1x2048xf32> loc(#loc5275)
    %22288 = stablehlo.add %21136, %22287 : tensor<1x2048xf32> loc(#loc5263)
    %22289 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22290 = stablehlo.broadcast_in_dim %22289, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22291 = stablehlo.multiply %22290, %arg154 : tensor<1x2048xf32> loc(#loc5275)
    %22292 = stablehlo.add %21156, %22291 : tensor<1x2048xf32> loc(#loc5263)
    %22293 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22294 = stablehlo.broadcast_in_dim %22293, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22295 = stablehlo.multiply %22294, %arg155 : tensor<1x2048xf32> loc(#loc5275)
    %22296 = stablehlo.add %21176, %22295 : tensor<1x2048xf32> loc(#loc5263)
    %22297 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22298 = stablehlo.broadcast_in_dim %22297, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22299 = stablehlo.multiply %22298, %arg156 : tensor<1x2048xf32> loc(#loc5275)
    %22300 = stablehlo.add %21196, %22299 : tensor<1x2048xf32> loc(#loc5263)
    %22301 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22302 = stablehlo.broadcast_in_dim %22301, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22303 = stablehlo.multiply %22302, %arg157 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22304 = stablehlo.add %21216, %22303 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22305 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22306 = stablehlo.broadcast_in_dim %22305, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22307 = stablehlo.multiply %22306, %arg158 : tensor<1x64xf32> loc(#loc5275)
    %22308 = stablehlo.add %21236, %22307 : tensor<1x64xf32> loc(#loc5263)
    %22309 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22310 = stablehlo.broadcast_in_dim %22309, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22311 = stablehlo.multiply %22310, %arg159 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22312 = stablehlo.add %21256, %22311 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22313 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22314 = stablehlo.broadcast_in_dim %22313, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22315 = stablehlo.multiply %22314, %arg160 : tensor<1x8192xf32> loc(#loc5275)
    %22316 = stablehlo.add %21276, %22315 : tensor<1x8192xf32> loc(#loc5263)
    %22317 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22318 = stablehlo.broadcast_in_dim %22317, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22319 = stablehlo.multiply %22318, %arg161 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22320 = stablehlo.add %21296, %22319 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22321 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22322 = stablehlo.broadcast_in_dim %22321, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22323 = stablehlo.multiply %22322, %arg162 : tensor<1x2048xf32> loc(#loc5275)
    %22324 = stablehlo.add %21316, %22323 : tensor<1x2048xf32> loc(#loc5263)
    %22325 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22326 = stablehlo.broadcast_in_dim %22325, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22327 = stablehlo.multiply %22326, %arg163 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22328 = stablehlo.add %21336, %22327 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22329 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22330 = stablehlo.broadcast_in_dim %22329, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22331 = stablehlo.multiply %22330, %arg164 : tensor<1x2048xf32> loc(#loc5275)
    %22332 = stablehlo.add %21356, %22331 : tensor<1x2048xf32> loc(#loc5263)
    %22333 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22334 = stablehlo.broadcast_in_dim %22333, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22335 = stablehlo.multiply %22334, %arg165 : tensor<1x2048xf32> loc(#loc5275)
    %22336 = stablehlo.add %21376, %22335 : tensor<1x2048xf32> loc(#loc5263)
    %22337 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22338 = stablehlo.broadcast_in_dim %22337, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22339 = stablehlo.multiply %22338, %arg166 : tensor<1x2048xf32> loc(#loc5275)
    %22340 = stablehlo.add %21396, %22339 : tensor<1x2048xf32> loc(#loc5263)
    %22341 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22342 = stablehlo.broadcast_in_dim %22341, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22343 = stablehlo.multiply %22342, %arg167 : tensor<1x2048xf32> loc(#loc5275)
    %22344 = stablehlo.add %21416, %22343 : tensor<1x2048xf32> loc(#loc5263)
    %22345 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22346 = stablehlo.broadcast_in_dim %22345, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22347 = stablehlo.multiply %22346, %arg168 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22348 = stablehlo.add %21436, %22347 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22349 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22350 = stablehlo.broadcast_in_dim %22349, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22351 = stablehlo.multiply %22350, %arg169 : tensor<1x64xf32> loc(#loc5275)
    %22352 = stablehlo.add %21456, %22351 : tensor<1x64xf32> loc(#loc5263)
    %22353 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22354 = stablehlo.broadcast_in_dim %22353, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22355 = stablehlo.multiply %22354, %arg170 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22356 = stablehlo.add %21476, %22355 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22357 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22358 = stablehlo.broadcast_in_dim %22357, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5275)
    %22359 = stablehlo.multiply %22358, %arg171 : tensor<1x8192xf32> loc(#loc5275)
    %22360 = stablehlo.add %21496, %22359 : tensor<1x8192xf32> loc(#loc5263)
    %22361 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22362 = stablehlo.broadcast_in_dim %22361, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5276)
    %22363 = stablehlo.multiply %22362, %arg172 : tensor<1x2048x8192xf32> loc(#loc5276)
    %22364 = stablehlo.add %21516, %22363 : tensor<1x2048x8192xf32> loc(#loc5263)
    %22365 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22366 = stablehlo.broadcast_in_dim %22365, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22367 = stablehlo.multiply %22366, %arg173 : tensor<1x2048xf32> loc(#loc5275)
    %22368 = stablehlo.add %21536, %22367 : tensor<1x2048xf32> loc(#loc5263)
    %22369 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22370 = stablehlo.broadcast_in_dim %22369, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5276)
    %22371 = stablehlo.multiply %22370, %arg174 : tensor<1x8192x2048xf32> loc(#loc5276)
    %22372 = stablehlo.add %21556, %22371 : tensor<1x8192x2048xf32> loc(#loc5263)
    %22373 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22374 = stablehlo.broadcast_in_dim %22373, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22375 = stablehlo.multiply %22374, %arg175 : tensor<1x2048xf32> loc(#loc5275)
    %22376 = stablehlo.add %21576, %22375 : tensor<1x2048xf32> loc(#loc5263)
    %22377 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22378 = stablehlo.broadcast_in_dim %22377, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22379 = stablehlo.multiply %22378, %arg176 : tensor<1x2048xf32> loc(#loc5275)
    %22380 = stablehlo.add %21596, %22379 : tensor<1x2048xf32> loc(#loc5263)
    %22381 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22382 = stablehlo.broadcast_in_dim %22381, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22383 = stablehlo.multiply %22382, %arg177 : tensor<1x2048xf32> loc(#loc5275)
    %22384 = stablehlo.add %21616, %22383 : tensor<1x2048xf32> loc(#loc5263)
    %22385 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22386 = stablehlo.broadcast_in_dim %22385, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5275)
    %22387 = stablehlo.multiply %22386, %arg178 : tensor<1x2048xf32> loc(#loc5275)
    %22388 = stablehlo.add %21636, %22387 : tensor<1x2048xf32> loc(#loc5263)
    %22389 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22390 = stablehlo.broadcast_in_dim %22389, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22391 = stablehlo.multiply %22390, %arg179 : tensor<1x3x2048x32x64xf32> loc(#loc5276)
    %22392 = stablehlo.add %21656, %22391 : tensor<1x3x2048x32x64xf32> loc(#loc5263)
    %22393 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22394 = stablehlo.broadcast_in_dim %22393, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5275)
    %22395 = stablehlo.multiply %22394, %arg180 : tensor<1x64xf32> loc(#loc5275)
    %22396 = stablehlo.add %21676, %22395 : tensor<1x64xf32> loc(#loc5263)
    %22397 = stablehlo.constant dense<1.000000e-03> : tensor<f32> loc(#loc)
    %22398 = stablehlo.broadcast_in_dim %22397, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5276)
    %22399 = stablehlo.multiply %22398, %arg181 : tensor<1x2048x32x64xf32> loc(#loc5276)
    %22400 = stablehlo.add %21696, %22399 : tensor<1x2048x32x64xf32> loc(#loc5263)
    %22401 = stablehlo.convert %arg198 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5225)
    %22402 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %22403 = stablehlo.broadcast_in_dim %22402, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5226)
    %22404 = stablehlo.subtract %22401, %22403 : tensor<1xf32> loc(#loc5226)
    %22405 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %22406 = stablehlo.broadcast_in_dim %22405, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5227)
    %22407 = stablehlo.divide %22404, %22406 : tensor<1xf32> loc(#loc5227)
    %22408 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %22409 = stablehlo.broadcast_in_dim %22408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5228)
    %22410 = stablehlo.multiply %22407, %22409 : tensor<1xf32> loc(#loc5228)
    %22411 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %22412 = stablehlo.broadcast_in_dim %22411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5229)
    %22413 = stablehlo.add %22412, %22410 : tensor<1xf32> loc(#loc5229)
    %22414 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %22415 = stablehlo.broadcast_in_dim %22414, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5230)
    %22416 = stablehlo.compare  LT, %22401, %22415,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc5230)
    %22417 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %22418 = stablehlo.broadcast_in_dim %22417, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5231)
    %22419 = stablehlo.compare  GE, %22401, %22418,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc5231)
    %22420 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %22421 = call @_where_2958(%22419, %22420, %22413) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc5232)
    %22422 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %22423 = call @_where_2959(%22416, %22422, %22421) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc5232)
    %22424 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %22425 = stablehlo.broadcast_in_dim %22424, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5277)
    %22426 = stablehlo.compare  LT, %arg198, %22425,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc5277)
    %22427 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %22428 = stablehlo.broadcast_in_dim %22427, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5234)
    %22429 = stablehlo.subtract %arg198, %22428 : tensor<1xi32> loc(#loc5234)
    %22430 = stablehlo.convert %22429 : (tensor<1xi32>) -> tensor<1xf32> loc(#loc5225)
    %22431 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %22432 = stablehlo.broadcast_in_dim %22431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5226)
    %22433 = stablehlo.subtract %22430, %22432 : tensor<1xf32> loc(#loc5226)
    %22434 = stablehlo.constant dense<4.999990e+05> : tensor<f32> loc(#loc)
    %22435 = stablehlo.broadcast_in_dim %22434, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5227)
    %22436 = stablehlo.divide %22433, %22435 : tensor<1xf32> loc(#loc5227)
    %22437 = stablehlo.constant dense<3.14159274> : tensor<f32> loc(#loc)
    %22438 = stablehlo.broadcast_in_dim %22437, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5228)
    %22439 = stablehlo.multiply %22436, %22438 : tensor<1xf32> loc(#loc5228)
    %22440 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %22441 = stablehlo.broadcast_in_dim %22440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5229)
    %22442 = stablehlo.add %22441, %22439 : tensor<1xf32> loc(#loc5229)
    %22443 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %22444 = stablehlo.broadcast_in_dim %22443, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5230)
    %22445 = stablehlo.compare  LT, %22430, %22444,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc5230)
    %22446 = stablehlo.constant dense<4.999990e+05> : tensor<f32> loc(#loc)
    %22447 = stablehlo.broadcast_in_dim %22446, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5231)
    %22448 = stablehlo.compare  GE, %22430, %22447,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc5231)
    %22449 = stablehlo.constant dense<3.14159274> : tensor<f32> loc(#loc)
    %22450 = call @_where_2960(%22448, %22449, %22442) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc5232)
    %22451 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %22452 = call @_where_2961(%22445, %22451, %22450) : (tensor<1xi1>, tensor<i32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc5232)
    %22453 = stablehlo.cosine %22452 : tensor<1xf32> loc(#loc5235)
    %22454 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %22455 = stablehlo.broadcast_in_dim %22454, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5229)
    %22456 = stablehlo.add %22455, %22453 : tensor<1xf32> loc(#loc5229)
    %22457 = stablehlo.constant dense<4.500000e-01> : tensor<f32> loc(#loc)
    %22458 = stablehlo.broadcast_in_dim %22457, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %22459 = stablehlo.multiply %22458, %22456 : tensor<1xf32> loc(#loc5236)
    %22460 = stablehlo.constant dense<1.000000e-01> : tensor<f32> loc(#loc)
    %22461 = stablehlo.broadcast_in_dim %22460, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5229)
    %22462 = stablehlo.add %22461, %22459 : tensor<1xf32> loc(#loc5229)
    %22463 = call @_where_2962(%22426, %22423, %22462) : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc5237)
    %22464 = stablehlo.constant dense<6.000000e-04> : tensor<f32> loc(#loc)
    %22465 = stablehlo.broadcast_in_dim %22464, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5228)
    %22466 = stablehlo.multiply %22463, %22465 : tensor<1xf32> loc(#loc5228)
    %22467 = stablehlo.constant dense<-1.000000e+00> : tensor<f32> loc(#loc)
    %22468 = stablehlo.broadcast_in_dim %22467, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5236)
    %22469 = stablehlo.multiply %22468, %22466 : tensor<1xf32> loc(#loc5236)
    %22470 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22471 = stablehlo.broadcast_in_dim %22470, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22472 = stablehlo.multiply %22471, %21700 : tensor<1x8192xf32> loc(#loc5262)
    %22473 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22474 = stablehlo.broadcast_in_dim %22473, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22475 = stablehlo.multiply %22474, %21704 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22476 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22477 = stablehlo.broadcast_in_dim %22476, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22478 = stablehlo.multiply %22477, %21708 : tensor<1x2048xf32> loc(#loc5262)
    %22479 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22480 = stablehlo.broadcast_in_dim %22479, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22481 = stablehlo.multiply %22480, %21712 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22482 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22483 = stablehlo.broadcast_in_dim %22482, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22484 = stablehlo.multiply %22483, %21716 : tensor<1x2048xf32> loc(#loc5262)
    %22485 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22486 = stablehlo.broadcast_in_dim %22485, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22487 = stablehlo.multiply %22486, %21720 : tensor<1x2048xf32> loc(#loc5262)
    %22488 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22489 = stablehlo.broadcast_in_dim %22488, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22490 = stablehlo.multiply %22489, %21724 : tensor<1x2048xf32> loc(#loc5262)
    %22491 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22492 = stablehlo.broadcast_in_dim %22491, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22493 = stablehlo.multiply %22492, %21728 : tensor<1x2048xf32> loc(#loc5262)
    %22494 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22495 = stablehlo.broadcast_in_dim %22494, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22496 = stablehlo.multiply %22495, %21732 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22497 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22498 = stablehlo.broadcast_in_dim %22497, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22499 = stablehlo.multiply %22498, %21736 : tensor<1x64xf32> loc(#loc5262)
    %22500 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22501 = stablehlo.broadcast_in_dim %22500, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22502 = stablehlo.multiply %22501, %21740 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22503 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22504 = stablehlo.broadcast_in_dim %22503, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22505 = stablehlo.multiply %22504, %21744 : tensor<1x8192xf32> loc(#loc5262)
    %22506 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22507 = stablehlo.broadcast_in_dim %22506, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22508 = stablehlo.multiply %22507, %21748 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22509 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22510 = stablehlo.broadcast_in_dim %22509, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22511 = stablehlo.multiply %22510, %21752 : tensor<1x2048xf32> loc(#loc5262)
    %22512 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22513 = stablehlo.broadcast_in_dim %22512, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22514 = stablehlo.multiply %22513, %21756 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22515 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22516 = stablehlo.broadcast_in_dim %22515, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22517 = stablehlo.multiply %22516, %21760 : tensor<1x2048xf32> loc(#loc5262)
    %22518 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22519 = stablehlo.broadcast_in_dim %22518, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22520 = stablehlo.multiply %22519, %21764 : tensor<1x2048xf32> loc(#loc5262)
    %22521 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22522 = stablehlo.broadcast_in_dim %22521, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22523 = stablehlo.multiply %22522, %21768 : tensor<1x2048xf32> loc(#loc5262)
    %22524 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22525 = stablehlo.broadcast_in_dim %22524, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22526 = stablehlo.multiply %22525, %21772 : tensor<1x2048xf32> loc(#loc5262)
    %22527 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22528 = stablehlo.broadcast_in_dim %22527, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22529 = stablehlo.multiply %22528, %21776 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22530 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22531 = stablehlo.broadcast_in_dim %22530, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22532 = stablehlo.multiply %22531, %21780 : tensor<1x64xf32> loc(#loc5262)
    %22533 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22534 = stablehlo.broadcast_in_dim %22533, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22535 = stablehlo.multiply %22534, %21784 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22536 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22537 = stablehlo.broadcast_in_dim %22536, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22538 = stablehlo.multiply %22537, %21788 : tensor<1x8192xf32> loc(#loc5262)
    %22539 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22540 = stablehlo.broadcast_in_dim %22539, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22541 = stablehlo.multiply %22540, %21792 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22542 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22543 = stablehlo.broadcast_in_dim %22542, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22544 = stablehlo.multiply %22543, %21796 : tensor<1x2048xf32> loc(#loc5262)
    %22545 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22546 = stablehlo.broadcast_in_dim %22545, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22547 = stablehlo.multiply %22546, %21800 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22548 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22549 = stablehlo.broadcast_in_dim %22548, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22550 = stablehlo.multiply %22549, %21804 : tensor<1x2048xf32> loc(#loc5262)
    %22551 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22552 = stablehlo.broadcast_in_dim %22551, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22553 = stablehlo.multiply %22552, %21808 : tensor<1x2048xf32> loc(#loc5262)
    %22554 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22555 = stablehlo.broadcast_in_dim %22554, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22556 = stablehlo.multiply %22555, %21812 : tensor<1x2048xf32> loc(#loc5262)
    %22557 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22558 = stablehlo.broadcast_in_dim %22557, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22559 = stablehlo.multiply %22558, %21816 : tensor<1x2048xf32> loc(#loc5262)
    %22560 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22561 = stablehlo.broadcast_in_dim %22560, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22562 = stablehlo.multiply %22561, %21820 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22563 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22564 = stablehlo.broadcast_in_dim %22563, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22565 = stablehlo.multiply %22564, %21824 : tensor<1x64xf32> loc(#loc5262)
    %22566 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22567 = stablehlo.broadcast_in_dim %22566, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22568 = stablehlo.multiply %22567, %21828 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22569 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22570 = stablehlo.broadcast_in_dim %22569, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22571 = stablehlo.multiply %22570, %21832 : tensor<1x8192xf32> loc(#loc5262)
    %22572 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22573 = stablehlo.broadcast_in_dim %22572, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22574 = stablehlo.multiply %22573, %21836 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22575 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22576 = stablehlo.broadcast_in_dim %22575, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22577 = stablehlo.multiply %22576, %21840 : tensor<1x2048xf32> loc(#loc5262)
    %22578 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22579 = stablehlo.broadcast_in_dim %22578, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22580 = stablehlo.multiply %22579, %21844 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22581 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22582 = stablehlo.broadcast_in_dim %22581, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22583 = stablehlo.multiply %22582, %21848 : tensor<1x2048xf32> loc(#loc5262)
    %22584 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22585 = stablehlo.broadcast_in_dim %22584, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22586 = stablehlo.multiply %22585, %21852 : tensor<1x2048xf32> loc(#loc5262)
    %22587 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22588 = stablehlo.broadcast_in_dim %22587, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22589 = stablehlo.multiply %22588, %21856 : tensor<1x2048xf32> loc(#loc5262)
    %22590 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22591 = stablehlo.broadcast_in_dim %22590, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22592 = stablehlo.multiply %22591, %21860 : tensor<1x2048xf32> loc(#loc5262)
    %22593 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22594 = stablehlo.broadcast_in_dim %22593, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22595 = stablehlo.multiply %22594, %21864 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22596 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22597 = stablehlo.broadcast_in_dim %22596, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22598 = stablehlo.multiply %22597, %21868 : tensor<1x64xf32> loc(#loc5262)
    %22599 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22600 = stablehlo.broadcast_in_dim %22599, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22601 = stablehlo.multiply %22600, %21872 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22602 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22603 = stablehlo.broadcast_in_dim %22602, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22604 = stablehlo.multiply %22603, %21876 : tensor<1x8192xf32> loc(#loc5262)
    %22605 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22606 = stablehlo.broadcast_in_dim %22605, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22607 = stablehlo.multiply %22606, %21880 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22608 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22609 = stablehlo.broadcast_in_dim %22608, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22610 = stablehlo.multiply %22609, %21884 : tensor<1x2048xf32> loc(#loc5262)
    %22611 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22612 = stablehlo.broadcast_in_dim %22611, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22613 = stablehlo.multiply %22612, %21888 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22614 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22615 = stablehlo.broadcast_in_dim %22614, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22616 = stablehlo.multiply %22615, %21892 : tensor<1x2048xf32> loc(#loc5262)
    %22617 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22618 = stablehlo.broadcast_in_dim %22617, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22619 = stablehlo.multiply %22618, %21896 : tensor<1x2048xf32> loc(#loc5262)
    %22620 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22621 = stablehlo.broadcast_in_dim %22620, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22622 = stablehlo.multiply %22621, %21900 : tensor<1x2048xf32> loc(#loc5262)
    %22623 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22624 = stablehlo.broadcast_in_dim %22623, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22625 = stablehlo.multiply %22624, %21904 : tensor<1x2048xf32> loc(#loc5262)
    %22626 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22627 = stablehlo.broadcast_in_dim %22626, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22628 = stablehlo.multiply %22627, %21908 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22629 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22630 = stablehlo.broadcast_in_dim %22629, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22631 = stablehlo.multiply %22630, %21912 : tensor<1x64xf32> loc(#loc5262)
    %22632 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22633 = stablehlo.broadcast_in_dim %22632, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22634 = stablehlo.multiply %22633, %21916 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22635 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22636 = stablehlo.broadcast_in_dim %22635, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22637 = stablehlo.multiply %22636, %21920 : tensor<1x8192xf32> loc(#loc5262)
    %22638 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22639 = stablehlo.broadcast_in_dim %22638, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22640 = stablehlo.multiply %22639, %21924 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22641 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22642 = stablehlo.broadcast_in_dim %22641, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22643 = stablehlo.multiply %22642, %21928 : tensor<1x2048xf32> loc(#loc5262)
    %22644 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22645 = stablehlo.broadcast_in_dim %22644, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22646 = stablehlo.multiply %22645, %21932 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22647 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22648 = stablehlo.broadcast_in_dim %22647, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22649 = stablehlo.multiply %22648, %21936 : tensor<1x2048xf32> loc(#loc5262)
    %22650 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22651 = stablehlo.broadcast_in_dim %22650, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22652 = stablehlo.multiply %22651, %21940 : tensor<1x2048xf32> loc(#loc5262)
    %22653 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22654 = stablehlo.broadcast_in_dim %22653, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22655 = stablehlo.multiply %22654, %21944 : tensor<1x2048xf32> loc(#loc5262)
    %22656 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22657 = stablehlo.broadcast_in_dim %22656, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22658 = stablehlo.multiply %22657, %21948 : tensor<1x2048xf32> loc(#loc5262)
    %22659 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22660 = stablehlo.broadcast_in_dim %22659, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22661 = stablehlo.multiply %22660, %21952 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22662 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22663 = stablehlo.broadcast_in_dim %22662, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22664 = stablehlo.multiply %22663, %21956 : tensor<1x64xf32> loc(#loc5262)
    %22665 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22666 = stablehlo.broadcast_in_dim %22665, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22667 = stablehlo.multiply %22666, %21960 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22668 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22669 = stablehlo.broadcast_in_dim %22668, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22670 = stablehlo.multiply %22669, %21964 : tensor<1x8192xf32> loc(#loc5262)
    %22671 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22672 = stablehlo.broadcast_in_dim %22671, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22673 = stablehlo.multiply %22672, %21968 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22674 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22675 = stablehlo.broadcast_in_dim %22674, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22676 = stablehlo.multiply %22675, %21972 : tensor<1x2048xf32> loc(#loc5262)
    %22677 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22678 = stablehlo.broadcast_in_dim %22677, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22679 = stablehlo.multiply %22678, %21976 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22680 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22681 = stablehlo.broadcast_in_dim %22680, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22682 = stablehlo.multiply %22681, %21980 : tensor<1x2048xf32> loc(#loc5262)
    %22683 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22684 = stablehlo.broadcast_in_dim %22683, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22685 = stablehlo.multiply %22684, %21984 : tensor<1x2048xf32> loc(#loc5262)
    %22686 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22687 = stablehlo.broadcast_in_dim %22686, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22688 = stablehlo.multiply %22687, %21988 : tensor<1x2048xf32> loc(#loc5262)
    %22689 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22690 = stablehlo.broadcast_in_dim %22689, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22691 = stablehlo.multiply %22690, %21992 : tensor<1x2048xf32> loc(#loc5262)
    %22692 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22693 = stablehlo.broadcast_in_dim %22692, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22694 = stablehlo.multiply %22693, %21996 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22695 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22696 = stablehlo.broadcast_in_dim %22695, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22697 = stablehlo.multiply %22696, %22000 : tensor<1x64xf32> loc(#loc5262)
    %22698 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22699 = stablehlo.broadcast_in_dim %22698, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22700 = stablehlo.multiply %22699, %22004 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22701 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22702 = stablehlo.broadcast_in_dim %22701, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22703 = stablehlo.multiply %22702, %22008 : tensor<1x8192xf32> loc(#loc5262)
    %22704 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22705 = stablehlo.broadcast_in_dim %22704, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22706 = stablehlo.multiply %22705, %22012 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22707 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22708 = stablehlo.broadcast_in_dim %22707, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22709 = stablehlo.multiply %22708, %22016 : tensor<1x2048xf32> loc(#loc5262)
    %22710 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22711 = stablehlo.broadcast_in_dim %22710, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22712 = stablehlo.multiply %22711, %22020 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22713 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22714 = stablehlo.broadcast_in_dim %22713, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22715 = stablehlo.multiply %22714, %22024 : tensor<1x2048xf32> loc(#loc5262)
    %22716 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22717 = stablehlo.broadcast_in_dim %22716, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22718 = stablehlo.multiply %22717, %22028 : tensor<1x2048xf32> loc(#loc5262)
    %22719 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22720 = stablehlo.broadcast_in_dim %22719, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22721 = stablehlo.multiply %22720, %22032 : tensor<1x2048xf32> loc(#loc5262)
    %22722 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22723 = stablehlo.broadcast_in_dim %22722, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22724 = stablehlo.multiply %22723, %22036 : tensor<1x2048xf32> loc(#loc5262)
    %22725 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22726 = stablehlo.broadcast_in_dim %22725, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22727 = stablehlo.multiply %22726, %22040 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22728 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22729 = stablehlo.broadcast_in_dim %22728, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22730 = stablehlo.multiply %22729, %22044 : tensor<1x64xf32> loc(#loc5262)
    %22731 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22732 = stablehlo.broadcast_in_dim %22731, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22733 = stablehlo.multiply %22732, %22048 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22734 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22735 = stablehlo.broadcast_in_dim %22734, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22736 = stablehlo.multiply %22735, %22052 : tensor<1x8192xf32> loc(#loc5262)
    %22737 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22738 = stablehlo.broadcast_in_dim %22737, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22739 = stablehlo.multiply %22738, %22056 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22740 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22741 = stablehlo.broadcast_in_dim %22740, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22742 = stablehlo.multiply %22741, %22060 : tensor<1x2048xf32> loc(#loc5262)
    %22743 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22744 = stablehlo.broadcast_in_dim %22743, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22745 = stablehlo.multiply %22744, %22064 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22746 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22747 = stablehlo.broadcast_in_dim %22746, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22748 = stablehlo.multiply %22747, %22068 : tensor<1x2048xf32> loc(#loc5262)
    %22749 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22750 = stablehlo.broadcast_in_dim %22749, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22751 = stablehlo.multiply %22750, %22072 : tensor<1x2048xf32> loc(#loc5262)
    %22752 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22753 = stablehlo.broadcast_in_dim %22752, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22754 = stablehlo.multiply %22753, %22076 : tensor<1x2048xf32> loc(#loc5262)
    %22755 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22756 = stablehlo.broadcast_in_dim %22755, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22757 = stablehlo.multiply %22756, %22080 : tensor<1x2048xf32> loc(#loc5262)
    %22758 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22759 = stablehlo.broadcast_in_dim %22758, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22760 = stablehlo.multiply %22759, %22084 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22761 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22762 = stablehlo.broadcast_in_dim %22761, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22763 = stablehlo.multiply %22762, %22088 : tensor<1x64xf32> loc(#loc5262)
    %22764 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22765 = stablehlo.broadcast_in_dim %22764, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22766 = stablehlo.multiply %22765, %22092 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22767 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22768 = stablehlo.broadcast_in_dim %22767, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22769 = stablehlo.multiply %22768, %22096 : tensor<1x8192xf32> loc(#loc5262)
    %22770 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22771 = stablehlo.broadcast_in_dim %22770, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22772 = stablehlo.multiply %22771, %22100 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22773 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22774 = stablehlo.broadcast_in_dim %22773, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22775 = stablehlo.multiply %22774, %22104 : tensor<1x2048xf32> loc(#loc5262)
    %22776 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22777 = stablehlo.broadcast_in_dim %22776, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22778 = stablehlo.multiply %22777, %22108 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22779 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22780 = stablehlo.broadcast_in_dim %22779, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22781 = stablehlo.multiply %22780, %22112 : tensor<1x2048xf32> loc(#loc5262)
    %22782 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22783 = stablehlo.broadcast_in_dim %22782, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22784 = stablehlo.multiply %22783, %22116 : tensor<1x2048xf32> loc(#loc5262)
    %22785 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22786 = stablehlo.broadcast_in_dim %22785, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22787 = stablehlo.multiply %22786, %22120 : tensor<1x2048xf32> loc(#loc5262)
    %22788 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22789 = stablehlo.broadcast_in_dim %22788, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22790 = stablehlo.multiply %22789, %22124 : tensor<1x2048xf32> loc(#loc5262)
    %22791 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22792 = stablehlo.broadcast_in_dim %22791, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22793 = stablehlo.multiply %22792, %22128 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22794 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22795 = stablehlo.broadcast_in_dim %22794, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22796 = stablehlo.multiply %22795, %22132 : tensor<1x64xf32> loc(#loc5262)
    %22797 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22798 = stablehlo.broadcast_in_dim %22797, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22799 = stablehlo.multiply %22798, %22136 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22800 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22801 = stablehlo.broadcast_in_dim %22800, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22802 = stablehlo.multiply %22801, %22140 : tensor<1x8192xf32> loc(#loc5262)
    %22803 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22804 = stablehlo.broadcast_in_dim %22803, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22805 = stablehlo.multiply %22804, %22144 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22806 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22807 = stablehlo.broadcast_in_dim %22806, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22808 = stablehlo.multiply %22807, %22148 : tensor<1x2048xf32> loc(#loc5262)
    %22809 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22810 = stablehlo.broadcast_in_dim %22809, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22811 = stablehlo.multiply %22810, %22152 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22812 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22813 = stablehlo.broadcast_in_dim %22812, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22814 = stablehlo.multiply %22813, %22156 : tensor<1x2048xf32> loc(#loc5262)
    %22815 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22816 = stablehlo.broadcast_in_dim %22815, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22817 = stablehlo.multiply %22816, %22160 : tensor<1x2048xf32> loc(#loc5262)
    %22818 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22819 = stablehlo.broadcast_in_dim %22818, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22820 = stablehlo.multiply %22819, %22164 : tensor<1x2048xf32> loc(#loc5262)
    %22821 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22822 = stablehlo.broadcast_in_dim %22821, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22823 = stablehlo.multiply %22822, %22168 : tensor<1x2048xf32> loc(#loc5262)
    %22824 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22825 = stablehlo.broadcast_in_dim %22824, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22826 = stablehlo.multiply %22825, %22172 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22827 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22828 = stablehlo.broadcast_in_dim %22827, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22829 = stablehlo.multiply %22828, %22176 : tensor<1x64xf32> loc(#loc5262)
    %22830 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22831 = stablehlo.broadcast_in_dim %22830, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22832 = stablehlo.multiply %22831, %22180 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22833 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22834 = stablehlo.broadcast_in_dim %22833, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22835 = stablehlo.multiply %22834, %22184 : tensor<1x8192xf32> loc(#loc5262)
    %22836 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22837 = stablehlo.broadcast_in_dim %22836, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22838 = stablehlo.multiply %22837, %22188 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22839 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22840 = stablehlo.broadcast_in_dim %22839, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22841 = stablehlo.multiply %22840, %22192 : tensor<1x2048xf32> loc(#loc5262)
    %22842 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22843 = stablehlo.broadcast_in_dim %22842, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22844 = stablehlo.multiply %22843, %22196 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22845 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22846 = stablehlo.broadcast_in_dim %22845, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22847 = stablehlo.multiply %22846, %22200 : tensor<1x2048xf32> loc(#loc5262)
    %22848 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22849 = stablehlo.broadcast_in_dim %22848, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22850 = stablehlo.multiply %22849, %22204 : tensor<1x2048xf32> loc(#loc5262)
    %22851 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22852 = stablehlo.broadcast_in_dim %22851, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22853 = stablehlo.multiply %22852, %22208 : tensor<1x2048xf32> loc(#loc5262)
    %22854 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22855 = stablehlo.broadcast_in_dim %22854, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22856 = stablehlo.multiply %22855, %22212 : tensor<1x2048xf32> loc(#loc5262)
    %22857 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22858 = stablehlo.broadcast_in_dim %22857, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22859 = stablehlo.multiply %22858, %22216 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22860 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22861 = stablehlo.broadcast_in_dim %22860, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22862 = stablehlo.multiply %22861, %22220 : tensor<1x64xf32> loc(#loc5262)
    %22863 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22864 = stablehlo.broadcast_in_dim %22863, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22865 = stablehlo.multiply %22864, %22224 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22866 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22867 = stablehlo.broadcast_in_dim %22866, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22868 = stablehlo.multiply %22867, %22228 : tensor<1x8192xf32> loc(#loc5262)
    %22869 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22870 = stablehlo.broadcast_in_dim %22869, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22871 = stablehlo.multiply %22870, %22232 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22872 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22873 = stablehlo.broadcast_in_dim %22872, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22874 = stablehlo.multiply %22873, %22236 : tensor<1x2048xf32> loc(#loc5262)
    %22875 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22876 = stablehlo.broadcast_in_dim %22875, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22877 = stablehlo.multiply %22876, %22240 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22878 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22879 = stablehlo.broadcast_in_dim %22878, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22880 = stablehlo.multiply %22879, %22244 : tensor<1x2048xf32> loc(#loc5262)
    %22881 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22882 = stablehlo.broadcast_in_dim %22881, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22883 = stablehlo.multiply %22882, %22248 : tensor<1x2048xf32> loc(#loc5262)
    %22884 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22885 = stablehlo.broadcast_in_dim %22884, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22886 = stablehlo.multiply %22885, %22252 : tensor<1x2048xf32> loc(#loc5262)
    %22887 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22888 = stablehlo.broadcast_in_dim %22887, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22889 = stablehlo.multiply %22888, %22256 : tensor<1x2048xf32> loc(#loc5262)
    %22890 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22891 = stablehlo.broadcast_in_dim %22890, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22892 = stablehlo.multiply %22891, %22260 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22893 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22894 = stablehlo.broadcast_in_dim %22893, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22895 = stablehlo.multiply %22894, %22264 : tensor<1x64xf32> loc(#loc5262)
    %22896 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22897 = stablehlo.broadcast_in_dim %22896, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22898 = stablehlo.multiply %22897, %22268 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22899 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22900 = stablehlo.broadcast_in_dim %22899, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22901 = stablehlo.multiply %22900, %22272 : tensor<1x8192xf32> loc(#loc5262)
    %22902 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22903 = stablehlo.broadcast_in_dim %22902, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22904 = stablehlo.multiply %22903, %22276 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22905 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22906 = stablehlo.broadcast_in_dim %22905, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22907 = stablehlo.multiply %22906, %22280 : tensor<1x2048xf32> loc(#loc5262)
    %22908 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22909 = stablehlo.broadcast_in_dim %22908, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22910 = stablehlo.multiply %22909, %22284 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22911 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22912 = stablehlo.broadcast_in_dim %22911, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22913 = stablehlo.multiply %22912, %22288 : tensor<1x2048xf32> loc(#loc5262)
    %22914 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22915 = stablehlo.broadcast_in_dim %22914, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22916 = stablehlo.multiply %22915, %22292 : tensor<1x2048xf32> loc(#loc5262)
    %22917 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22918 = stablehlo.broadcast_in_dim %22917, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22919 = stablehlo.multiply %22918, %22296 : tensor<1x2048xf32> loc(#loc5262)
    %22920 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22921 = stablehlo.broadcast_in_dim %22920, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22922 = stablehlo.multiply %22921, %22300 : tensor<1x2048xf32> loc(#loc5262)
    %22923 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22924 = stablehlo.broadcast_in_dim %22923, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22925 = stablehlo.multiply %22924, %22304 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22926 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22927 = stablehlo.broadcast_in_dim %22926, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22928 = stablehlo.multiply %22927, %22308 : tensor<1x64xf32> loc(#loc5262)
    %22929 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22930 = stablehlo.broadcast_in_dim %22929, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22931 = stablehlo.multiply %22930, %22312 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22932 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22933 = stablehlo.broadcast_in_dim %22932, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22934 = stablehlo.multiply %22933, %22316 : tensor<1x8192xf32> loc(#loc5262)
    %22935 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22936 = stablehlo.broadcast_in_dim %22935, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22937 = stablehlo.multiply %22936, %22320 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22938 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22939 = stablehlo.broadcast_in_dim %22938, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22940 = stablehlo.multiply %22939, %22324 : tensor<1x2048xf32> loc(#loc5262)
    %22941 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22942 = stablehlo.broadcast_in_dim %22941, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22943 = stablehlo.multiply %22942, %22328 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22944 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22945 = stablehlo.broadcast_in_dim %22944, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22946 = stablehlo.multiply %22945, %22332 : tensor<1x2048xf32> loc(#loc5262)
    %22947 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22948 = stablehlo.broadcast_in_dim %22947, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22949 = stablehlo.multiply %22948, %22336 : tensor<1x2048xf32> loc(#loc5262)
    %22950 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22951 = stablehlo.broadcast_in_dim %22950, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22952 = stablehlo.multiply %22951, %22340 : tensor<1x2048xf32> loc(#loc5262)
    %22953 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22954 = stablehlo.broadcast_in_dim %22953, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22955 = stablehlo.multiply %22954, %22344 : tensor<1x2048xf32> loc(#loc5262)
    %22956 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22957 = stablehlo.broadcast_in_dim %22956, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22958 = stablehlo.multiply %22957, %22348 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22959 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22960 = stablehlo.broadcast_in_dim %22959, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22961 = stablehlo.multiply %22960, %22352 : tensor<1x64xf32> loc(#loc5262)
    %22962 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22963 = stablehlo.broadcast_in_dim %22962, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22964 = stablehlo.multiply %22963, %22356 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22965 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22966 = stablehlo.broadcast_in_dim %22965, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x8192xf32> loc(#loc5262)
    %22967 = stablehlo.multiply %22966, %22360 : tensor<1x8192xf32> loc(#loc5262)
    %22968 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22969 = stablehlo.broadcast_in_dim %22968, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x2048x8192xf32> loc(#loc5262)
    %22970 = stablehlo.multiply %22969, %22364 : tensor<1x2048x8192xf32> loc(#loc5262)
    %22971 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22972 = stablehlo.broadcast_in_dim %22971, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22973 = stablehlo.multiply %22972, %22368 : tensor<1x2048xf32> loc(#loc5262)
    %22974 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1xf32> loc(#loc5279)
    %22975 = stablehlo.broadcast_in_dim %22974, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<1x8192x2048xf32> loc(#loc5262)
    %22976 = stablehlo.multiply %22975, %22372 : tensor<1x8192x2048xf32> loc(#loc5262)
    %22977 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22978 = stablehlo.broadcast_in_dim %22977, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22979 = stablehlo.multiply %22978, %22376 : tensor<1x2048xf32> loc(#loc5262)
    %22980 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22981 = stablehlo.broadcast_in_dim %22980, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22982 = stablehlo.multiply %22981, %22380 : tensor<1x2048xf32> loc(#loc5262)
    %22983 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22984 = stablehlo.broadcast_in_dim %22983, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22985 = stablehlo.multiply %22984, %22384 : tensor<1x2048xf32> loc(#loc5262)
    %22986 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22987 = stablehlo.broadcast_in_dim %22986, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x2048xf32> loc(#loc5262)
    %22988 = stablehlo.multiply %22987, %22388 : tensor<1x2048xf32> loc(#loc5262)
    %22989 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1x1xf32> loc(#loc5280)
    %22990 = stablehlo.broadcast_in_dim %22989, dims = [0, 1, 2, 3, 4] : (tensor<1x1x1x1x1xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22991 = stablehlo.multiply %22990, %22392 : tensor<1x3x2048x32x64xf32> loc(#loc5262)
    %22992 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc5278)
    %22993 = stablehlo.broadcast_in_dim %22992, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<1x64xf32> loc(#loc5262)
    %22994 = stablehlo.multiply %22993, %22396 : tensor<1x64xf32> loc(#loc5262)
    %22995 = stablehlo.broadcast_in_dim %22469, dims = [0] : (tensor<1xf32>) -> tensor<1x1x1x1xf32> loc(#loc5281)
    %22996 = stablehlo.broadcast_in_dim %22995, dims = [0, 1, 2, 3] : (tensor<1x1x1x1xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5262)
    %22997 = stablehlo.multiply %22996, %22400 : tensor<1x2048x32x64xf32> loc(#loc5262)
    %22998 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %22999 = stablehlo.broadcast_in_dim %22998, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5256)
    %23000 = stablehlo.add %arg198, %22999 : tensor<1xi32> loc(#loc5256)
    %23001 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %23002 = stablehlo.broadcast_in_dim %23001, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5256)
    %23003 = stablehlo.add %arg551, %23002 : tensor<1xi32> loc(#loc5256)
    %23004 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23005 = stablehlo.broadcast_in_dim %23004, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5290)
    %23006 = call @_where_2963(%2694, %4254, %23005) : (tensor<i1>, tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc5291)
    %23007 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23008 = stablehlo.broadcast_in_dim %23007, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5290)
    %23009 = call @_where_2964(%2694, %4256, %23008) : (tensor<i1>, tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc5291)
    %23010 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23011 = stablehlo.broadcast_in_dim %23010, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5292)
    %23012 = call @_where_2965(%2694, %4258, %23011) : (tensor<i1>, tensor<2048x2048xf32>, tensor<2048x2048xf32>) -> tensor<2048x2048xf32> loc(#loc5291)
    %23013 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23014 = stablehlo.broadcast_in_dim %23013, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5293)
    %23015 = call @_where_2966(%2694, %4260, %23014) : (tensor<i1>, tensor<51200xf32>, tensor<51200xf32>) -> tensor<51200xf32> loc(#loc5291)
    %23016 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23017 = stablehlo.broadcast_in_dim %23016, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5294)
    %23018 = call @_where_2967(%2694, %4262, %23017) : (tensor<i1>, tensor<2048x51200xf32>, tensor<2048x51200xf32>) -> tensor<2048x51200xf32> loc(#loc5291)
    %23019 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23020 = stablehlo.broadcast_in_dim %23019, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23021 = call @_where_2968(%2694, %22472, %23020) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23022 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23023 = stablehlo.broadcast_in_dim %23022, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23024 = call @_where_2969(%2694, %22475, %23023) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23025 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23026 = stablehlo.broadcast_in_dim %23025, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23027 = call @_where_2970(%2694, %22478, %23026) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23028 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23029 = stablehlo.broadcast_in_dim %23028, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23030 = call @_where_2971(%2694, %22481, %23029) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23031 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23032 = stablehlo.broadcast_in_dim %23031, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23033 = call @_where_2972(%2694, %22484, %23032) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23034 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23035 = stablehlo.broadcast_in_dim %23034, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23036 = call @_where_2973(%2694, %22487, %23035) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23037 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23038 = stablehlo.broadcast_in_dim %23037, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23039 = call @_where_2974(%2694, %22490, %23038) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23040 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23041 = stablehlo.broadcast_in_dim %23040, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23042 = call @_where_2975(%2694, %22493, %23041) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23043 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23044 = stablehlo.broadcast_in_dim %23043, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23045 = call @_where_2976(%2694, %22496, %23044) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23046 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23047 = stablehlo.broadcast_in_dim %23046, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23048 = call @_where_2977(%2694, %22499, %23047) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23049 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23050 = stablehlo.broadcast_in_dim %23049, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23051 = call @_where_2978(%2694, %22502, %23050) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23052 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23053 = stablehlo.broadcast_in_dim %23052, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23054 = call @_where_2979(%2694, %22505, %23053) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23055 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23056 = stablehlo.broadcast_in_dim %23055, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23057 = call @_where_2980(%2694, %22508, %23056) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23058 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23059 = stablehlo.broadcast_in_dim %23058, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23060 = call @_where_2981(%2694, %22511, %23059) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23061 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23062 = stablehlo.broadcast_in_dim %23061, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23063 = call @_where_2982(%2694, %22514, %23062) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23064 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23065 = stablehlo.broadcast_in_dim %23064, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23066 = call @_where_2983(%2694, %22517, %23065) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23067 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23068 = stablehlo.broadcast_in_dim %23067, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23069 = call @_where_2984(%2694, %22520, %23068) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23070 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23071 = stablehlo.broadcast_in_dim %23070, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23072 = call @_where_2985(%2694, %22523, %23071) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23073 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23074 = stablehlo.broadcast_in_dim %23073, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23075 = call @_where_2986(%2694, %22526, %23074) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23076 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23077 = stablehlo.broadcast_in_dim %23076, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23078 = call @_where_2987(%2694, %22529, %23077) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23079 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23080 = stablehlo.broadcast_in_dim %23079, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23081 = call @_where_2988(%2694, %22532, %23080) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23082 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23083 = stablehlo.broadcast_in_dim %23082, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23084 = call @_where_2989(%2694, %22535, %23083) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23085 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23086 = stablehlo.broadcast_in_dim %23085, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23087 = call @_where_2990(%2694, %22538, %23086) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23088 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23089 = stablehlo.broadcast_in_dim %23088, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23090 = call @_where_2991(%2694, %22541, %23089) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23091 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23092 = stablehlo.broadcast_in_dim %23091, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23093 = call @_where_2992(%2694, %22544, %23092) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23094 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23095 = stablehlo.broadcast_in_dim %23094, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23096 = call @_where_2993(%2694, %22547, %23095) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23097 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23098 = stablehlo.broadcast_in_dim %23097, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23099 = call @_where_2994(%2694, %22550, %23098) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23100 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23101 = stablehlo.broadcast_in_dim %23100, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23102 = call @_where_2995(%2694, %22553, %23101) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23103 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23104 = stablehlo.broadcast_in_dim %23103, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23105 = call @_where_2996(%2694, %22556, %23104) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23106 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23107 = stablehlo.broadcast_in_dim %23106, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23108 = call @_where_2997(%2694, %22559, %23107) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23109 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23110 = stablehlo.broadcast_in_dim %23109, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23111 = call @_where_2998(%2694, %22562, %23110) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23112 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23113 = stablehlo.broadcast_in_dim %23112, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23114 = call @_where_2999(%2694, %22565, %23113) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23115 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23116 = stablehlo.broadcast_in_dim %23115, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23117 = call @_where_3000(%2694, %22568, %23116) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23118 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23119 = stablehlo.broadcast_in_dim %23118, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23120 = call @_where_3001(%2694, %22571, %23119) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23121 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23122 = stablehlo.broadcast_in_dim %23121, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23123 = call @_where_3002(%2694, %22574, %23122) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23124 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23125 = stablehlo.broadcast_in_dim %23124, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23126 = call @_where_3003(%2694, %22577, %23125) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23127 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23128 = stablehlo.broadcast_in_dim %23127, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23129 = call @_where_3004(%2694, %22580, %23128) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23130 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23131 = stablehlo.broadcast_in_dim %23130, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23132 = call @_where_3005(%2694, %22583, %23131) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23133 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23134 = stablehlo.broadcast_in_dim %23133, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23135 = call @_where_3006(%2694, %22586, %23134) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23136 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23137 = stablehlo.broadcast_in_dim %23136, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23138 = call @_where_3007(%2694, %22589, %23137) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23139 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23140 = stablehlo.broadcast_in_dim %23139, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23141 = call @_where_3008(%2694, %22592, %23140) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23142 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23143 = stablehlo.broadcast_in_dim %23142, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23144 = call @_where_3009(%2694, %22595, %23143) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23145 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23146 = stablehlo.broadcast_in_dim %23145, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23147 = call @_where_3010(%2694, %22598, %23146) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23148 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23149 = stablehlo.broadcast_in_dim %23148, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23150 = call @_where_3011(%2694, %22601, %23149) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23151 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23152 = stablehlo.broadcast_in_dim %23151, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23153 = call @_where_3012(%2694, %22604, %23152) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23154 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23155 = stablehlo.broadcast_in_dim %23154, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23156 = call @_where_3013(%2694, %22607, %23155) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23157 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23158 = stablehlo.broadcast_in_dim %23157, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23159 = call @_where_3014(%2694, %22610, %23158) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23160 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23161 = stablehlo.broadcast_in_dim %23160, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23162 = call @_where_3015(%2694, %22613, %23161) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23163 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23164 = stablehlo.broadcast_in_dim %23163, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23165 = call @_where_3016(%2694, %22616, %23164) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23166 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23167 = stablehlo.broadcast_in_dim %23166, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23168 = call @_where_3017(%2694, %22619, %23167) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23169 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23170 = stablehlo.broadcast_in_dim %23169, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23171 = call @_where_3018(%2694, %22622, %23170) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23172 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23173 = stablehlo.broadcast_in_dim %23172, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23174 = call @_where_3019(%2694, %22625, %23173) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23175 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23176 = stablehlo.broadcast_in_dim %23175, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23177 = call @_where_3020(%2694, %22628, %23176) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23178 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23179 = stablehlo.broadcast_in_dim %23178, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23180 = call @_where_3021(%2694, %22631, %23179) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23181 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23182 = stablehlo.broadcast_in_dim %23181, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23183 = call @_where_3022(%2694, %22634, %23182) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23184 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23185 = stablehlo.broadcast_in_dim %23184, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23186 = call @_where_3023(%2694, %22637, %23185) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23187 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23188 = stablehlo.broadcast_in_dim %23187, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23189 = call @_where_3024(%2694, %22640, %23188) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23190 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23191 = stablehlo.broadcast_in_dim %23190, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23192 = call @_where_3025(%2694, %22643, %23191) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23193 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23194 = stablehlo.broadcast_in_dim %23193, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23195 = call @_where_3026(%2694, %22646, %23194) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23196 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23197 = stablehlo.broadcast_in_dim %23196, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23198 = call @_where_3027(%2694, %22649, %23197) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23199 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23200 = stablehlo.broadcast_in_dim %23199, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23201 = call @_where_3028(%2694, %22652, %23200) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23202 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23203 = stablehlo.broadcast_in_dim %23202, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23204 = call @_where_3029(%2694, %22655, %23203) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23205 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23206 = stablehlo.broadcast_in_dim %23205, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23207 = call @_where_3030(%2694, %22658, %23206) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23208 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23209 = stablehlo.broadcast_in_dim %23208, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23210 = call @_where_3031(%2694, %22661, %23209) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23211 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23212 = stablehlo.broadcast_in_dim %23211, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23213 = call @_where_3032(%2694, %22664, %23212) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23214 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23215 = stablehlo.broadcast_in_dim %23214, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23216 = call @_where_3033(%2694, %22667, %23215) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23217 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23218 = stablehlo.broadcast_in_dim %23217, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23219 = call @_where_3034(%2694, %22670, %23218) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23220 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23221 = stablehlo.broadcast_in_dim %23220, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23222 = call @_where_3035(%2694, %22673, %23221) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23223 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23224 = stablehlo.broadcast_in_dim %23223, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23225 = call @_where_3036(%2694, %22676, %23224) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23226 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23227 = stablehlo.broadcast_in_dim %23226, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23228 = call @_where_3037(%2694, %22679, %23227) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23229 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23230 = stablehlo.broadcast_in_dim %23229, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23231 = call @_where_3038(%2694, %22682, %23230) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23232 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23233 = stablehlo.broadcast_in_dim %23232, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23234 = call @_where_3039(%2694, %22685, %23233) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23235 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23236 = stablehlo.broadcast_in_dim %23235, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23237 = call @_where_3040(%2694, %22688, %23236) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23238 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23239 = stablehlo.broadcast_in_dim %23238, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23240 = call @_where_3041(%2694, %22691, %23239) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23241 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23242 = stablehlo.broadcast_in_dim %23241, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23243 = call @_where_3042(%2694, %22694, %23242) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23244 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23245 = stablehlo.broadcast_in_dim %23244, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23246 = call @_where_3043(%2694, %22697, %23245) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23247 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23248 = stablehlo.broadcast_in_dim %23247, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23249 = call @_where_3044(%2694, %22700, %23248) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23250 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23251 = stablehlo.broadcast_in_dim %23250, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23252 = call @_where_3045(%2694, %22703, %23251) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23253 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23254 = stablehlo.broadcast_in_dim %23253, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23255 = call @_where_3046(%2694, %22706, %23254) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23256 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23257 = stablehlo.broadcast_in_dim %23256, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23258 = call @_where_3047(%2694, %22709, %23257) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23259 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23260 = stablehlo.broadcast_in_dim %23259, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23261 = call @_where_3048(%2694, %22712, %23260) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23262 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23263 = stablehlo.broadcast_in_dim %23262, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23264 = call @_where_3049(%2694, %22715, %23263) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23265 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23266 = stablehlo.broadcast_in_dim %23265, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23267 = call @_where_3050(%2694, %22718, %23266) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23268 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23269 = stablehlo.broadcast_in_dim %23268, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23270 = call @_where_3051(%2694, %22721, %23269) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23271 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23272 = stablehlo.broadcast_in_dim %23271, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23273 = call @_where_3052(%2694, %22724, %23272) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23274 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23275 = stablehlo.broadcast_in_dim %23274, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23276 = call @_where_3053(%2694, %22727, %23275) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23277 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23278 = stablehlo.broadcast_in_dim %23277, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23279 = call @_where_3054(%2694, %22730, %23278) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23280 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23281 = stablehlo.broadcast_in_dim %23280, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23282 = call @_where_3055(%2694, %22733, %23281) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23283 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23284 = stablehlo.broadcast_in_dim %23283, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23285 = call @_where_3056(%2694, %22736, %23284) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23286 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23287 = stablehlo.broadcast_in_dim %23286, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23288 = call @_where_3057(%2694, %22739, %23287) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23289 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23290 = stablehlo.broadcast_in_dim %23289, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23291 = call @_where_3058(%2694, %22742, %23290) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23292 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23293 = stablehlo.broadcast_in_dim %23292, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23294 = call @_where_3059(%2694, %22745, %23293) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23295 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23296 = stablehlo.broadcast_in_dim %23295, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23297 = call @_where_3060(%2694, %22748, %23296) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23298 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23299 = stablehlo.broadcast_in_dim %23298, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23300 = call @_where_3061(%2694, %22751, %23299) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23301 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23302 = stablehlo.broadcast_in_dim %23301, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23303 = call @_where_3062(%2694, %22754, %23302) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23304 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23305 = stablehlo.broadcast_in_dim %23304, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23306 = call @_where_3063(%2694, %22757, %23305) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23307 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23308 = stablehlo.broadcast_in_dim %23307, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23309 = call @_where_3064(%2694, %22760, %23308) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23310 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23311 = stablehlo.broadcast_in_dim %23310, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23312 = call @_where_3065(%2694, %22763, %23311) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23313 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23314 = stablehlo.broadcast_in_dim %23313, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23315 = call @_where_3066(%2694, %22766, %23314) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23316 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23317 = stablehlo.broadcast_in_dim %23316, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23318 = call @_where_3067(%2694, %22769, %23317) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23319 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23320 = stablehlo.broadcast_in_dim %23319, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23321 = call @_where_3068(%2694, %22772, %23320) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23322 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23323 = stablehlo.broadcast_in_dim %23322, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23324 = call @_where_3069(%2694, %22775, %23323) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23325 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23326 = stablehlo.broadcast_in_dim %23325, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23327 = call @_where_3070(%2694, %22778, %23326) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23328 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23329 = stablehlo.broadcast_in_dim %23328, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23330 = call @_where_3071(%2694, %22781, %23329) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23331 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23332 = stablehlo.broadcast_in_dim %23331, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23333 = call @_where_3072(%2694, %22784, %23332) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23334 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23335 = stablehlo.broadcast_in_dim %23334, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23336 = call @_where_3073(%2694, %22787, %23335) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23337 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23338 = stablehlo.broadcast_in_dim %23337, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23339 = call @_where_3074(%2694, %22790, %23338) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23340 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23341 = stablehlo.broadcast_in_dim %23340, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23342 = call @_where_3075(%2694, %22793, %23341) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23343 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23344 = stablehlo.broadcast_in_dim %23343, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23345 = call @_where_3076(%2694, %22796, %23344) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23346 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23347 = stablehlo.broadcast_in_dim %23346, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23348 = call @_where_3077(%2694, %22799, %23347) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23350 = stablehlo.broadcast_in_dim %23349, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23351 = call @_where_3078(%2694, %22802, %23350) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23352 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23353 = stablehlo.broadcast_in_dim %23352, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23354 = call @_where_3079(%2694, %22805, %23353) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23355 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23356 = stablehlo.broadcast_in_dim %23355, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23357 = call @_where_3080(%2694, %22808, %23356) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23358 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23359 = stablehlo.broadcast_in_dim %23358, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23360 = call @_where_3081(%2694, %22811, %23359) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23361 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23362 = stablehlo.broadcast_in_dim %23361, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23363 = call @_where_3082(%2694, %22814, %23362) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23364 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23365 = stablehlo.broadcast_in_dim %23364, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23366 = call @_where_3083(%2694, %22817, %23365) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23367 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23368 = stablehlo.broadcast_in_dim %23367, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23369 = call @_where_3084(%2694, %22820, %23368) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23370 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23371 = stablehlo.broadcast_in_dim %23370, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23372 = call @_where_3085(%2694, %22823, %23371) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23373 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23374 = stablehlo.broadcast_in_dim %23373, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23375 = call @_where_3086(%2694, %22826, %23374) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23376 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23377 = stablehlo.broadcast_in_dim %23376, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23378 = call @_where_3087(%2694, %22829, %23377) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23379 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23380 = stablehlo.broadcast_in_dim %23379, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23381 = call @_where_3088(%2694, %22832, %23380) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23382 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23383 = stablehlo.broadcast_in_dim %23382, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23384 = call @_where_3089(%2694, %22835, %23383) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23385 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23386 = stablehlo.broadcast_in_dim %23385, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23387 = call @_where_3090(%2694, %22838, %23386) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23388 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23389 = stablehlo.broadcast_in_dim %23388, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23390 = call @_where_3091(%2694, %22841, %23389) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23391 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23392 = stablehlo.broadcast_in_dim %23391, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23393 = call @_where_3092(%2694, %22844, %23392) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23394 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23395 = stablehlo.broadcast_in_dim %23394, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23396 = call @_where_3093(%2694, %22847, %23395) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23397 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23398 = stablehlo.broadcast_in_dim %23397, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23399 = call @_where_3094(%2694, %22850, %23398) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23400 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23401 = stablehlo.broadcast_in_dim %23400, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23402 = call @_where_3095(%2694, %22853, %23401) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23403 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23404 = stablehlo.broadcast_in_dim %23403, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23405 = call @_where_3096(%2694, %22856, %23404) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23406 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23407 = stablehlo.broadcast_in_dim %23406, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23408 = call @_where_3097(%2694, %22859, %23407) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23409 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23410 = stablehlo.broadcast_in_dim %23409, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23411 = call @_where_3098(%2694, %22862, %23410) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23412 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23413 = stablehlo.broadcast_in_dim %23412, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23414 = call @_where_3099(%2694, %22865, %23413) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23415 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23416 = stablehlo.broadcast_in_dim %23415, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23417 = call @_where_3100(%2694, %22868, %23416) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23418 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23419 = stablehlo.broadcast_in_dim %23418, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23420 = call @_where_3101(%2694, %22871, %23419) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23421 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23422 = stablehlo.broadcast_in_dim %23421, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23423 = call @_where_3102(%2694, %22874, %23422) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23424 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23425 = stablehlo.broadcast_in_dim %23424, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23426 = call @_where_3103(%2694, %22877, %23425) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23427 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23428 = stablehlo.broadcast_in_dim %23427, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23429 = call @_where_3104(%2694, %22880, %23428) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23430 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23431 = stablehlo.broadcast_in_dim %23430, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23432 = call @_where_3105(%2694, %22883, %23431) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23433 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23434 = stablehlo.broadcast_in_dim %23433, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23435 = call @_where_3106(%2694, %22886, %23434) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23436 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23437 = stablehlo.broadcast_in_dim %23436, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23438 = call @_where_3107(%2694, %22889, %23437) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23439 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23440 = stablehlo.broadcast_in_dim %23439, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23441 = call @_where_3108(%2694, %22892, %23440) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23442 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23443 = stablehlo.broadcast_in_dim %23442, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23444 = call @_where_3109(%2694, %22895, %23443) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23445 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23446 = stablehlo.broadcast_in_dim %23445, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23447 = call @_where_3110(%2694, %22898, %23446) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23448 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23449 = stablehlo.broadcast_in_dim %23448, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23450 = call @_where_3111(%2694, %22901, %23449) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23451 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23452 = stablehlo.broadcast_in_dim %23451, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23453 = call @_where_3112(%2694, %22904, %23452) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23454 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23455 = stablehlo.broadcast_in_dim %23454, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23456 = call @_where_3113(%2694, %22907, %23455) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23457 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23458 = stablehlo.broadcast_in_dim %23457, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23459 = call @_where_3114(%2694, %22910, %23458) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23460 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23461 = stablehlo.broadcast_in_dim %23460, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23462 = call @_where_3115(%2694, %22913, %23461) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23463 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23464 = stablehlo.broadcast_in_dim %23463, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23465 = call @_where_3116(%2694, %22916, %23464) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23466 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23467 = stablehlo.broadcast_in_dim %23466, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23468 = call @_where_3117(%2694, %22919, %23467) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23469 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23470 = stablehlo.broadcast_in_dim %23469, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23471 = call @_where_3118(%2694, %22922, %23470) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23472 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23473 = stablehlo.broadcast_in_dim %23472, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23474 = call @_where_3119(%2694, %22925, %23473) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23475 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23476 = stablehlo.broadcast_in_dim %23475, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23477 = call @_where_3120(%2694, %22928, %23476) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23478 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23479 = stablehlo.broadcast_in_dim %23478, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23480 = call @_where_3121(%2694, %22931, %23479) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23481 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23482 = stablehlo.broadcast_in_dim %23481, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23483 = call @_where_3122(%2694, %22934, %23482) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23484 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23485 = stablehlo.broadcast_in_dim %23484, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23486 = call @_where_3123(%2694, %22937, %23485) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23487 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23488 = stablehlo.broadcast_in_dim %23487, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23489 = call @_where_3124(%2694, %22940, %23488) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23490 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23491 = stablehlo.broadcast_in_dim %23490, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23492 = call @_where_3125(%2694, %22943, %23491) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23493 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23494 = stablehlo.broadcast_in_dim %23493, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23495 = call @_where_3126(%2694, %22946, %23494) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23496 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23497 = stablehlo.broadcast_in_dim %23496, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23498 = call @_where_3127(%2694, %22949, %23497) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23499 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23500 = stablehlo.broadcast_in_dim %23499, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23501 = call @_where_3128(%2694, %22952, %23500) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23502 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23503 = stablehlo.broadcast_in_dim %23502, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23504 = call @_where_3129(%2694, %22955, %23503) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23505 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23506 = stablehlo.broadcast_in_dim %23505, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23507 = call @_where_3130(%2694, %22958, %23506) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23508 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23509 = stablehlo.broadcast_in_dim %23508, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23510 = call @_where_3131(%2694, %22961, %23509) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23511 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23512 = stablehlo.broadcast_in_dim %23511, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23513 = call @_where_3132(%2694, %22964, %23512) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23514 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23515 = stablehlo.broadcast_in_dim %23514, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5295)
    %23516 = call @_where_3133(%2694, %22967, %23515) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5291)
    %23517 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23518 = stablehlo.broadcast_in_dim %23517, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5296)
    %23519 = call @_where_3134(%2694, %22970, %23518) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5291)
    %23520 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23521 = stablehlo.broadcast_in_dim %23520, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23522 = call @_where_3135(%2694, %22973, %23521) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23523 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23524 = stablehlo.broadcast_in_dim %23523, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5298)
    %23525 = call @_where_3136(%2694, %22976, %23524) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5291)
    %23526 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23527 = stablehlo.broadcast_in_dim %23526, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23528 = call @_where_3137(%2694, %22979, %23527) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23529 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23530 = stablehlo.broadcast_in_dim %23529, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23531 = call @_where_3138(%2694, %22982, %23530) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23532 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23533 = stablehlo.broadcast_in_dim %23532, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23534 = call @_where_3139(%2694, %22985, %23533) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23535 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23536 = stablehlo.broadcast_in_dim %23535, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5297)
    %23537 = call @_where_3140(%2694, %22988, %23536) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5291)
    %23538 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23539 = stablehlo.broadcast_in_dim %23538, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5299)
    %23540 = call @_where_3141(%2694, %22991, %23539) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5291)
    %23541 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23542 = stablehlo.broadcast_in_dim %23541, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5300)
    %23543 = call @_where_3142(%2694, %22994, %23542) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5291)
    %23544 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc)
    %23545 = stablehlo.broadcast_in_dim %23544, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5301)
    %23546 = call @_where_3143(%2694, %22997, %23545) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5291)
    %23547 = call @_where_3144(%2694, %3803, %arg182) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5302)
    %23548 = call @_where_3145(%2694, %3805, %arg183) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5302)
    %23549 = call @_where_3146(%2694, %4264, %arg184) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5302)
    %23550 = call @_where_3147(%2694, %3864, %arg185) : (tensor<i1>, tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc5302)
    %23551 = call @_where_3148(%2694, %3911, %arg186) : (tensor<i1>, tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc5302)
    %23552 = call @_where_3149(%2694, %3958, %arg187) : (tensor<i1>, tensor<2048x2048xf32>, tensor<2048x2048xf32>) -> tensor<2048x2048xf32> loc(#loc5302)
    %23553 = call @_where_3150(%2694, %4005, %arg188) : (tensor<i1>, tensor<51200xf32>, tensor<51200xf32>) -> tensor<51200xf32> loc(#loc5302)
    %23554 = call @_where_3151(%2694, %4052, %arg189) : (tensor<i1>, tensor<2048x51200xf32>, tensor<2048x51200xf32>) -> tensor<2048x51200xf32> loc(#loc5302)
    %23555 = call @_where_3152(%2694, %3872, %arg190) : (tensor<i1>, tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc5302)
    %23556 = call @_where_3153(%2694, %3919, %arg191) : (tensor<i1>, tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc5302)
    %23557 = call @_where_3154(%2694, %3966, %arg192) : (tensor<i1>, tensor<2048x2048xf32>, tensor<2048x2048xf32>) -> tensor<2048x2048xf32> loc(#loc5302)
    %23558 = call @_where_3155(%2694, %4013, %arg193) : (tensor<i1>, tensor<51200xf32>, tensor<51200xf32>) -> tensor<51200xf32> loc(#loc5302)
    %23559 = call @_where_3156(%2694, %4060, %arg194) : (tensor<i1>, tensor<2048x51200xf32>, tensor<2048x51200xf32>) -> tensor<2048x51200xf32> loc(#loc5302)
    %23560 = call @_where_3157(%2694, %4266, %arg195) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5302)
    %23561 = call @_where_3158(%2694, %4269, %arg196) : (tensor<i1>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> loc(#loc5302)
    %23562 = call @_where_3159(%2694, %4272, %arg197) : (tensor<i1>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> loc(#loc5302)
    %23563 = call @_where_3160(%2694, %23000, %arg198) : (tensor<i1>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> loc(#loc5302)
    %23564 = call @_where_3161(%2694, %5032, %arg199) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23565 = call @_where_3162(%2694, %5099, %arg200) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23566 = call @_where_3163(%2694, %5166, %arg201) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23567 = call @_where_3164(%2694, %5233, %arg202) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23568 = call @_where_3165(%2694, %5300, %arg203) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23569 = call @_where_3166(%2694, %5367, %arg204) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23570 = call @_where_3167(%2694, %5434, %arg205) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23571 = call @_where_3168(%2694, %5501, %arg206) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23572 = call @_where_3169(%2694, %5568, %arg207) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23573 = call @_where_3170(%2694, %5635, %arg208) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23574 = call @_where_3171(%2694, %5702, %arg209) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23575 = call @_where_3172(%2694, %5769, %arg210) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23576 = call @_where_3173(%2694, %5836, %arg211) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23577 = call @_where_3174(%2694, %5903, %arg212) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23578 = call @_where_3175(%2694, %5970, %arg213) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23579 = call @_where_3176(%2694, %6037, %arg214) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23580 = call @_where_3177(%2694, %6104, %arg215) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23581 = call @_where_3178(%2694, %6171, %arg216) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23582 = call @_where_3179(%2694, %6238, %arg217) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23583 = call @_where_3180(%2694, %6305, %arg218) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23584 = call @_where_3181(%2694, %6372, %arg219) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23585 = call @_where_3182(%2694, %6439, %arg220) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23586 = call @_where_3183(%2694, %6506, %arg221) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23587 = call @_where_3184(%2694, %6573, %arg222) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23588 = call @_where_3185(%2694, %6640, %arg223) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23589 = call @_where_3186(%2694, %6707, %arg224) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23590 = call @_where_3187(%2694, %6774, %arg225) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23591 = call @_where_3188(%2694, %6841, %arg226) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23592 = call @_where_3189(%2694, %6908, %arg227) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23593 = call @_where_3190(%2694, %6975, %arg228) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23594 = call @_where_3191(%2694, %7042, %arg229) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23595 = call @_where_3192(%2694, %7109, %arg230) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23596 = call @_where_3193(%2694, %7176, %arg231) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23597 = call @_where_3194(%2694, %7243, %arg232) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23598 = call @_where_3195(%2694, %7310, %arg233) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23599 = call @_where_3196(%2694, %7377, %arg234) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23600 = call @_where_3197(%2694, %7444, %arg235) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23601 = call @_where_3198(%2694, %7511, %arg236) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23602 = call @_where_3199(%2694, %7578, %arg237) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23603 = call @_where_3200(%2694, %7645, %arg238) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23604 = call @_where_3201(%2694, %7712, %arg239) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23605 = call @_where_3202(%2694, %7779, %arg240) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23606 = call @_where_3203(%2694, %7846, %arg241) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23607 = call @_where_3204(%2694, %7913, %arg242) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23608 = call @_where_3205(%2694, %7980, %arg243) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23609 = call @_where_3206(%2694, %8047, %arg244) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23610 = call @_where_3207(%2694, %8114, %arg245) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23611 = call @_where_3208(%2694, %8181, %arg246) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23612 = call @_where_3209(%2694, %8248, %arg247) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23613 = call @_where_3210(%2694, %8315, %arg248) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23614 = call @_where_3211(%2694, %8382, %arg249) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23615 = call @_where_3212(%2694, %8449, %arg250) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23616 = call @_where_3213(%2694, %8516, %arg251) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23617 = call @_where_3214(%2694, %8583, %arg252) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23618 = call @_where_3215(%2694, %8650, %arg253) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23619 = call @_where_3216(%2694, %8717, %arg254) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23620 = call @_where_3217(%2694, %8784, %arg255) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23621 = call @_where_3218(%2694, %8851, %arg256) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23622 = call @_where_3219(%2694, %8918, %arg257) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23623 = call @_where_3220(%2694, %8985, %arg258) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23624 = call @_where_3221(%2694, %9052, %arg259) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23625 = call @_where_3222(%2694, %9119, %arg260) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23626 = call @_where_3223(%2694, %9186, %arg261) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23627 = call @_where_3224(%2694, %9253, %arg262) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23628 = call @_where_3225(%2694, %9320, %arg263) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23629 = call @_where_3226(%2694, %9387, %arg264) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23630 = call @_where_3227(%2694, %9454, %arg265) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23631 = call @_where_3228(%2694, %9521, %arg266) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23632 = call @_where_3229(%2694, %9588, %arg267) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23633 = call @_where_3230(%2694, %9655, %arg268) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23634 = call @_where_3231(%2694, %9722, %arg269) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23635 = call @_where_3232(%2694, %9789, %arg270) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23636 = call @_where_3233(%2694, %9856, %arg271) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23637 = call @_where_3234(%2694, %9923, %arg272) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23638 = call @_where_3235(%2694, %9990, %arg273) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23639 = call @_where_3236(%2694, %10057, %arg274) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23640 = call @_where_3237(%2694, %10124, %arg275) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23641 = call @_where_3238(%2694, %10191, %arg276) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23642 = call @_where_3239(%2694, %10258, %arg277) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23643 = call @_where_3240(%2694, %10325, %arg278) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23644 = call @_where_3241(%2694, %10392, %arg279) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23645 = call @_where_3242(%2694, %10459, %arg280) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23646 = call @_where_3243(%2694, %10526, %arg281) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23647 = call @_where_3244(%2694, %10593, %arg282) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23648 = call @_where_3245(%2694, %10660, %arg283) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23649 = call @_where_3246(%2694, %10727, %arg284) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23650 = call @_where_3247(%2694, %10794, %arg285) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23651 = call @_where_3248(%2694, %10861, %arg286) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23652 = call @_where_3249(%2694, %10928, %arg287) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23653 = call @_where_3250(%2694, %10995, %arg288) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23654 = call @_where_3251(%2694, %11062, %arg289) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23655 = call @_where_3252(%2694, %11129, %arg290) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23656 = call @_where_3253(%2694, %11196, %arg291) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23657 = call @_where_3254(%2694, %11263, %arg292) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23658 = call @_where_3255(%2694, %11330, %arg293) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23659 = call @_where_3256(%2694, %11397, %arg294) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23660 = call @_where_3257(%2694, %11464, %arg295) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23661 = call @_where_3258(%2694, %11531, %arg296) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23662 = call @_where_3259(%2694, %11598, %arg297) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23663 = call @_where_3260(%2694, %11665, %arg298) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23664 = call @_where_3261(%2694, %11732, %arg299) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23665 = call @_where_3262(%2694, %11799, %arg300) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23666 = call @_where_3263(%2694, %11866, %arg301) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23667 = call @_where_3264(%2694, %11933, %arg302) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23668 = call @_where_3265(%2694, %12000, %arg303) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23669 = call @_where_3266(%2694, %12067, %arg304) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23670 = call @_where_3267(%2694, %12134, %arg305) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23671 = call @_where_3268(%2694, %12201, %arg306) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23672 = call @_where_3269(%2694, %12268, %arg307) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23673 = call @_where_3270(%2694, %12335, %arg308) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23674 = call @_where_3271(%2694, %12402, %arg309) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23675 = call @_where_3272(%2694, %12469, %arg310) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23676 = call @_where_3273(%2694, %12536, %arg311) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23677 = call @_where_3274(%2694, %12603, %arg312) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23678 = call @_where_3275(%2694, %12670, %arg313) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23679 = call @_where_3276(%2694, %12737, %arg314) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23680 = call @_where_3277(%2694, %12804, %arg315) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23681 = call @_where_3278(%2694, %12871, %arg316) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23682 = call @_where_3279(%2694, %12938, %arg317) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23683 = call @_where_3280(%2694, %13005, %arg318) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23684 = call @_where_3281(%2694, %13072, %arg319) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23685 = call @_where_3282(%2694, %13139, %arg320) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23686 = call @_where_3283(%2694, %13206, %arg321) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23687 = call @_where_3284(%2694, %13273, %arg322) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23688 = call @_where_3285(%2694, %13340, %arg323) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23689 = call @_where_3286(%2694, %13407, %arg324) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23690 = call @_where_3287(%2694, %13474, %arg325) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23691 = call @_where_3288(%2694, %13541, %arg326) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23692 = call @_where_3289(%2694, %13608, %arg327) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23693 = call @_where_3290(%2694, %13675, %arg328) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23694 = call @_where_3291(%2694, %13742, %arg329) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23695 = call @_where_3292(%2694, %13809, %arg330) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23696 = call @_where_3293(%2694, %13876, %arg331) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23697 = call @_where_3294(%2694, %13943, %arg332) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23698 = call @_where_3295(%2694, %14010, %arg333) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23699 = call @_where_3296(%2694, %14077, %arg334) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23700 = call @_where_3297(%2694, %14144, %arg335) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23701 = call @_where_3298(%2694, %14211, %arg336) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23702 = call @_where_3299(%2694, %14278, %arg337) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23703 = call @_where_3300(%2694, %14345, %arg338) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23704 = call @_where_3301(%2694, %14412, %arg339) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23705 = call @_where_3302(%2694, %14479, %arg340) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23706 = call @_where_3303(%2694, %14546, %arg341) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23707 = call @_where_3304(%2694, %14613, %arg342) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23708 = call @_where_3305(%2694, %14680, %arg343) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23709 = call @_where_3306(%2694, %14747, %arg344) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23710 = call @_where_3307(%2694, %14814, %arg345) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23711 = call @_where_3308(%2694, %14881, %arg346) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23712 = call @_where_3309(%2694, %14948, %arg347) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23713 = call @_where_3310(%2694, %15015, %arg348) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23714 = call @_where_3311(%2694, %15082, %arg349) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23715 = call @_where_3312(%2694, %15149, %arg350) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23716 = call @_where_3313(%2694, %15216, %arg351) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23717 = call @_where_3314(%2694, %15283, %arg352) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23718 = call @_where_3315(%2694, %15350, %arg353) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23719 = call @_where_3316(%2694, %15417, %arg354) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23720 = call @_where_3317(%2694, %15484, %arg355) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23721 = call @_where_3318(%2694, %15551, %arg356) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23722 = call @_where_3319(%2694, %15618, %arg357) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23723 = call @_where_3320(%2694, %15685, %arg358) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23724 = call @_where_3321(%2694, %15752, %arg359) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23725 = call @_where_3322(%2694, %15819, %arg360) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23726 = call @_where_3323(%2694, %15886, %arg361) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23727 = call @_where_3324(%2694, %15953, %arg362) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23728 = call @_where_3325(%2694, %16020, %arg363) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23729 = call @_where_3326(%2694, %16087, %arg364) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23730 = call @_where_3327(%2694, %16154, %arg365) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23731 = call @_where_3328(%2694, %16221, %arg366) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23732 = call @_where_3329(%2694, %16288, %arg367) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23733 = call @_where_3330(%2694, %16355, %arg368) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23734 = call @_where_3331(%2694, %16422, %arg369) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23735 = call @_where_3332(%2694, %16489, %arg370) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23736 = call @_where_3333(%2694, %16556, %arg371) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23737 = call @_where_3334(%2694, %16623, %arg372) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23738 = call @_where_3335(%2694, %16690, %arg373) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23739 = call @_where_3336(%2694, %16757, %arg374) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23740 = call @_where_3337(%2694, %5043, %arg375) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23741 = call @_where_3338(%2694, %5110, %arg376) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23742 = call @_where_3339(%2694, %5177, %arg377) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23743 = call @_where_3340(%2694, %5244, %arg378) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23744 = call @_where_3341(%2694, %5311, %arg379) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23745 = call @_where_3342(%2694, %5378, %arg380) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23746 = call @_where_3343(%2694, %5445, %arg381) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23747 = call @_where_3344(%2694, %5512, %arg382) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23748 = call @_where_3345(%2694, %5579, %arg383) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23749 = call @_where_3346(%2694, %5646, %arg384) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23750 = call @_where_3347(%2694, %5713, %arg385) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23751 = call @_where_3348(%2694, %5780, %arg386) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23752 = call @_where_3349(%2694, %5847, %arg387) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23753 = call @_where_3350(%2694, %5914, %arg388) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23754 = call @_where_3351(%2694, %5981, %arg389) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23755 = call @_where_3352(%2694, %6048, %arg390) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23756 = call @_where_3353(%2694, %6115, %arg391) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23757 = call @_where_3354(%2694, %6182, %arg392) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23758 = call @_where_3355(%2694, %6249, %arg393) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23759 = call @_where_3356(%2694, %6316, %arg394) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23760 = call @_where_3357(%2694, %6383, %arg395) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23761 = call @_where_3358(%2694, %6450, %arg396) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23762 = call @_where_3359(%2694, %6517, %arg397) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23763 = call @_where_3360(%2694, %6584, %arg398) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23764 = call @_where_3361(%2694, %6651, %arg399) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23765 = call @_where_3362(%2694, %6718, %arg400) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23766 = call @_where_3363(%2694, %6785, %arg401) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23767 = call @_where_3364(%2694, %6852, %arg402) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23768 = call @_where_3365(%2694, %6919, %arg403) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23769 = call @_where_3366(%2694, %6986, %arg404) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23770 = call @_where_3367(%2694, %7053, %arg405) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23771 = call @_where_3368(%2694, %7120, %arg406) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23772 = call @_where_3369(%2694, %7187, %arg407) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23773 = call @_where_3370(%2694, %7254, %arg408) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23774 = call @_where_3371(%2694, %7321, %arg409) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23775 = call @_where_3372(%2694, %7388, %arg410) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23776 = call @_where_3373(%2694, %7455, %arg411) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23777 = call @_where_3374(%2694, %7522, %arg412) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23778 = call @_where_3375(%2694, %7589, %arg413) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23779 = call @_where_3376(%2694, %7656, %arg414) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23780 = call @_where_3377(%2694, %7723, %arg415) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23781 = call @_where_3378(%2694, %7790, %arg416) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23782 = call @_where_3379(%2694, %7857, %arg417) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23783 = call @_where_3380(%2694, %7924, %arg418) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23784 = call @_where_3381(%2694, %7991, %arg419) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23785 = call @_where_3382(%2694, %8058, %arg420) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23786 = call @_where_3383(%2694, %8125, %arg421) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23787 = call @_where_3384(%2694, %8192, %arg422) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23788 = call @_where_3385(%2694, %8259, %arg423) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23789 = call @_where_3386(%2694, %8326, %arg424) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23790 = call @_where_3387(%2694, %8393, %arg425) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23791 = call @_where_3388(%2694, %8460, %arg426) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23792 = call @_where_3389(%2694, %8527, %arg427) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23793 = call @_where_3390(%2694, %8594, %arg428) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23794 = call @_where_3391(%2694, %8661, %arg429) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23795 = call @_where_3392(%2694, %8728, %arg430) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23796 = call @_where_3393(%2694, %8795, %arg431) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23797 = call @_where_3394(%2694, %8862, %arg432) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23798 = call @_where_3395(%2694, %8929, %arg433) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23799 = call @_where_3396(%2694, %8996, %arg434) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23800 = call @_where_3397(%2694, %9063, %arg435) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23801 = call @_where_3398(%2694, %9130, %arg436) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23802 = call @_where_3399(%2694, %9197, %arg437) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23803 = call @_where_3400(%2694, %9264, %arg438) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23804 = call @_where_3401(%2694, %9331, %arg439) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23805 = call @_where_3402(%2694, %9398, %arg440) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23806 = call @_where_3403(%2694, %9465, %arg441) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23807 = call @_where_3404(%2694, %9532, %arg442) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23808 = call @_where_3405(%2694, %9599, %arg443) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23809 = call @_where_3406(%2694, %9666, %arg444) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23810 = call @_where_3407(%2694, %9733, %arg445) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23811 = call @_where_3408(%2694, %9800, %arg446) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23812 = call @_where_3409(%2694, %9867, %arg447) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23813 = call @_where_3410(%2694, %9934, %arg448) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23814 = call @_where_3411(%2694, %10001, %arg449) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23815 = call @_where_3412(%2694, %10068, %arg450) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23816 = call @_where_3413(%2694, %10135, %arg451) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23817 = call @_where_3414(%2694, %10202, %arg452) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23818 = call @_where_3415(%2694, %10269, %arg453) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23819 = call @_where_3416(%2694, %10336, %arg454) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23820 = call @_where_3417(%2694, %10403, %arg455) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23821 = call @_where_3418(%2694, %10470, %arg456) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23822 = call @_where_3419(%2694, %10537, %arg457) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23823 = call @_where_3420(%2694, %10604, %arg458) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23824 = call @_where_3421(%2694, %10671, %arg459) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23825 = call @_where_3422(%2694, %10738, %arg460) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23826 = call @_where_3423(%2694, %10805, %arg461) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23827 = call @_where_3424(%2694, %10872, %arg462) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23828 = call @_where_3425(%2694, %10939, %arg463) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23829 = call @_where_3426(%2694, %11006, %arg464) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23830 = call @_where_3427(%2694, %11073, %arg465) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23831 = call @_where_3428(%2694, %11140, %arg466) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23832 = call @_where_3429(%2694, %11207, %arg467) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23833 = call @_where_3430(%2694, %11274, %arg468) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23834 = call @_where_3431(%2694, %11341, %arg469) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23835 = call @_where_3432(%2694, %11408, %arg470) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23836 = call @_where_3433(%2694, %11475, %arg471) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23837 = call @_where_3434(%2694, %11542, %arg472) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23838 = call @_where_3435(%2694, %11609, %arg473) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23839 = call @_where_3436(%2694, %11676, %arg474) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23840 = call @_where_3437(%2694, %11743, %arg475) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23841 = call @_where_3438(%2694, %11810, %arg476) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23842 = call @_where_3439(%2694, %11877, %arg477) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23843 = call @_where_3440(%2694, %11944, %arg478) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23844 = call @_where_3441(%2694, %12011, %arg479) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23845 = call @_where_3442(%2694, %12078, %arg480) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23846 = call @_where_3443(%2694, %12145, %arg481) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23847 = call @_where_3444(%2694, %12212, %arg482) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23848 = call @_where_3445(%2694, %12279, %arg483) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23849 = call @_where_3446(%2694, %12346, %arg484) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23850 = call @_where_3447(%2694, %12413, %arg485) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23851 = call @_where_3448(%2694, %12480, %arg486) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23852 = call @_where_3449(%2694, %12547, %arg487) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23853 = call @_where_3450(%2694, %12614, %arg488) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23854 = call @_where_3451(%2694, %12681, %arg489) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23855 = call @_where_3452(%2694, %12748, %arg490) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23856 = call @_where_3453(%2694, %12815, %arg491) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23857 = call @_where_3454(%2694, %12882, %arg492) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23858 = call @_where_3455(%2694, %12949, %arg493) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23859 = call @_where_3456(%2694, %13016, %arg494) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23860 = call @_where_3457(%2694, %13083, %arg495) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23861 = call @_where_3458(%2694, %13150, %arg496) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23862 = call @_where_3459(%2694, %13217, %arg497) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23863 = call @_where_3460(%2694, %13284, %arg498) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23864 = call @_where_3461(%2694, %13351, %arg499) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23865 = call @_where_3462(%2694, %13418, %arg500) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23866 = call @_where_3463(%2694, %13485, %arg501) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23867 = call @_where_3464(%2694, %13552, %arg502) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23868 = call @_where_3465(%2694, %13619, %arg503) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23869 = call @_where_3466(%2694, %13686, %arg504) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23870 = call @_where_3467(%2694, %13753, %arg505) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23871 = call @_where_3468(%2694, %13820, %arg506) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23872 = call @_where_3469(%2694, %13887, %arg507) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23873 = call @_where_3470(%2694, %13954, %arg508) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23874 = call @_where_3471(%2694, %14021, %arg509) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23875 = call @_where_3472(%2694, %14088, %arg510) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23876 = call @_where_3473(%2694, %14155, %arg511) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23877 = call @_where_3474(%2694, %14222, %arg512) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23878 = call @_where_3475(%2694, %14289, %arg513) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23879 = call @_where_3476(%2694, %14356, %arg514) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23880 = call @_where_3477(%2694, %14423, %arg515) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23881 = call @_where_3478(%2694, %14490, %arg516) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23882 = call @_where_3479(%2694, %14557, %arg517) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23883 = call @_where_3480(%2694, %14624, %arg518) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23884 = call @_where_3481(%2694, %14691, %arg519) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23885 = call @_where_3482(%2694, %14758, %arg520) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23886 = call @_where_3483(%2694, %14825, %arg521) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23887 = call @_where_3484(%2694, %14892, %arg522) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23888 = call @_where_3485(%2694, %14959, %arg523) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23889 = call @_where_3486(%2694, %15026, %arg524) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23890 = call @_where_3487(%2694, %15093, %arg525) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23891 = call @_where_3488(%2694, %15160, %arg526) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23892 = call @_where_3489(%2694, %15227, %arg527) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23893 = call @_where_3490(%2694, %15294, %arg528) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23894 = call @_where_3491(%2694, %15361, %arg529) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23895 = call @_where_3492(%2694, %15428, %arg530) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23896 = call @_where_3493(%2694, %15495, %arg531) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23897 = call @_where_3494(%2694, %15562, %arg532) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23898 = call @_where_3495(%2694, %15629, %arg533) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23899 = call @_where_3496(%2694, %15696, %arg534) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23900 = call @_where_3497(%2694, %15763, %arg535) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23901 = call @_where_3498(%2694, %15830, %arg536) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23902 = call @_where_3499(%2694, %15897, %arg537) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23903 = call @_where_3500(%2694, %15964, %arg538) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23904 = call @_where_3501(%2694, %16031, %arg539) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23905 = call @_where_3502(%2694, %16098, %arg540) : (tensor<i1>, tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc5302)
    %23906 = call @_where_3503(%2694, %16165, %arg541) : (tensor<i1>, tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc5302)
    %23907 = call @_where_3504(%2694, %16232, %arg542) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23908 = call @_where_3505(%2694, %16299, %arg543) : (tensor<i1>, tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc5302)
    %23909 = call @_where_3506(%2694, %16366, %arg544) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23910 = call @_where_3507(%2694, %16433, %arg545) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23911 = call @_where_3508(%2694, %16500, %arg546) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23912 = call @_where_3509(%2694, %16567, %arg547) : (tensor<i1>, tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc5302)
    %23913 = call @_where_3510(%2694, %16634, %arg548) : (tensor<i1>, tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5302)
    %23914 = call @_where_3511(%2694, %16701, %arg549) : (tensor<i1>, tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc5302)
    %23915 = call @_where_3512(%2694, %16768, %arg550) : (tensor<i1>, tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc5302)
    %23916 = call @_where_3513(%2694, %23003, %arg551) : (tensor<i1>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> loc(#loc5302)
    %23917 = stablehlo.multiply %23006, %23006 : tensor<2048xf32> loc(#loc5238)
    %23918 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %23919 = stablehlo.reduce(%23917 init: %23918) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %23920 = stablehlo.multiply %23009, %23009 : tensor<2048xf32> loc(#loc5238)
    %23921 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %23922 = stablehlo.reduce(%23920 init: %23921) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %23923 = stablehlo.multiply %23012, %23012 : tensor<2048x2048xf32> loc(#loc5238)
    %23924 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %23925 = stablehlo.reduce(%23923 init: %23924) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %23926 = stablehlo.multiply %23015, %23015 : tensor<51200xf32> loc(#loc5238)
    %23927 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %23928 = stablehlo.reduce(%23926 init: %23927) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %23929 = stablehlo.multiply %23018, %23018 : tensor<2048x51200xf32> loc(#loc5238)
    %23930 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %23931 = stablehlo.reduce(%23929 init: %23930) across dimensions = [0, 1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %23932 = stablehlo.multiply %23021, %23021 : tensor<1x8192xf32> loc(#loc5238)
    %23933 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %23934 = stablehlo.reduce(%23932 init: %23933) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %23935 = stablehlo.multiply %23024, %23024 : tensor<1x2048x8192xf32> loc(#loc5238)
    %23936 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %23937 = stablehlo.reduce(%23935 init: %23936) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %23938 = stablehlo.multiply %23027, %23027 : tensor<1x2048xf32> loc(#loc5242)
    %23939 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23940 = stablehlo.reduce(%23938 init: %23939) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23941 = stablehlo.multiply %23030, %23030 : tensor<1x8192x2048xf32> loc(#loc5238)
    %23942 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %23943 = stablehlo.reduce(%23941 init: %23942) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %23944 = stablehlo.multiply %23033, %23033 : tensor<1x2048xf32> loc(#loc5242)
    %23945 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23946 = stablehlo.reduce(%23944 init: %23945) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23947 = stablehlo.multiply %23036, %23036 : tensor<1x2048xf32> loc(#loc5242)
    %23948 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23949 = stablehlo.reduce(%23947 init: %23948) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23950 = stablehlo.multiply %23039, %23039 : tensor<1x2048xf32> loc(#loc5242)
    %23951 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23952 = stablehlo.reduce(%23950 init: %23951) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23953 = stablehlo.multiply %23042, %23042 : tensor<1x2048xf32> loc(#loc5242)
    %23954 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23955 = stablehlo.reduce(%23953 init: %23954) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23956 = stablehlo.multiply %23045, %23045 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %23957 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %23958 = stablehlo.reduce(%23956 init: %23957) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %23959 = stablehlo.multiply %23048, %23048 : tensor<1x64xf32> loc(#loc5238)
    %23960 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %23961 = stablehlo.reduce(%23959 init: %23960) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %23962 = stablehlo.multiply %23051, %23051 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %23963 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %23964 = stablehlo.reduce(%23962 init: %23963) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %23965 = stablehlo.multiply %23054, %23054 : tensor<1x8192xf32> loc(#loc5238)
    %23966 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %23967 = stablehlo.reduce(%23965 init: %23966) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %23968 = stablehlo.multiply %23057, %23057 : tensor<1x2048x8192xf32> loc(#loc5238)
    %23969 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %23970 = stablehlo.reduce(%23968 init: %23969) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %23971 = stablehlo.multiply %23060, %23060 : tensor<1x2048xf32> loc(#loc5242)
    %23972 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23973 = stablehlo.reduce(%23971 init: %23972) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23974 = stablehlo.multiply %23063, %23063 : tensor<1x8192x2048xf32> loc(#loc5238)
    %23975 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %23976 = stablehlo.reduce(%23974 init: %23975) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %23977 = stablehlo.multiply %23066, %23066 : tensor<1x2048xf32> loc(#loc5242)
    %23978 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23979 = stablehlo.reduce(%23977 init: %23978) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23980 = stablehlo.multiply %23069, %23069 : tensor<1x2048xf32> loc(#loc5242)
    %23981 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23982 = stablehlo.reduce(%23980 init: %23981) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23983 = stablehlo.multiply %23072, %23072 : tensor<1x2048xf32> loc(#loc5242)
    %23984 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23985 = stablehlo.reduce(%23983 init: %23984) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23986 = stablehlo.multiply %23075, %23075 : tensor<1x2048xf32> loc(#loc5242)
    %23987 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %23988 = stablehlo.reduce(%23986 init: %23987) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %23989 = stablehlo.multiply %23078, %23078 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %23990 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %23991 = stablehlo.reduce(%23989 init: %23990) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %23992 = stablehlo.multiply %23081, %23081 : tensor<1x64xf32> loc(#loc5238)
    %23993 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %23994 = stablehlo.reduce(%23992 init: %23993) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %23995 = stablehlo.multiply %23084, %23084 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %23996 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %23997 = stablehlo.reduce(%23995 init: %23996) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %23998 = stablehlo.multiply %23087, %23087 : tensor<1x8192xf32> loc(#loc5238)
    %23999 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24000 = stablehlo.reduce(%23998 init: %23999) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24001 = stablehlo.multiply %23090, %23090 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24002 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24003 = stablehlo.reduce(%24001 init: %24002) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24004 = stablehlo.multiply %23093, %23093 : tensor<1x2048xf32> loc(#loc5242)
    %24005 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24006 = stablehlo.reduce(%24004 init: %24005) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24007 = stablehlo.multiply %23096, %23096 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24008 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24009 = stablehlo.reduce(%24007 init: %24008) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24010 = stablehlo.multiply %23099, %23099 : tensor<1x2048xf32> loc(#loc5242)
    %24011 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24012 = stablehlo.reduce(%24010 init: %24011) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24013 = stablehlo.multiply %23102, %23102 : tensor<1x2048xf32> loc(#loc5242)
    %24014 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24015 = stablehlo.reduce(%24013 init: %24014) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24016 = stablehlo.multiply %23105, %23105 : tensor<1x2048xf32> loc(#loc5242)
    %24017 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24018 = stablehlo.reduce(%24016 init: %24017) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24019 = stablehlo.multiply %23108, %23108 : tensor<1x2048xf32> loc(#loc5242)
    %24020 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24021 = stablehlo.reduce(%24019 init: %24020) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24022 = stablehlo.multiply %23111, %23111 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24023 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24024 = stablehlo.reduce(%24022 init: %24023) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24025 = stablehlo.multiply %23114, %23114 : tensor<1x64xf32> loc(#loc5238)
    %24026 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24027 = stablehlo.reduce(%24025 init: %24026) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24028 = stablehlo.multiply %23117, %23117 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24029 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24030 = stablehlo.reduce(%24028 init: %24029) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24031 = stablehlo.multiply %23120, %23120 : tensor<1x8192xf32> loc(#loc5238)
    %24032 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24033 = stablehlo.reduce(%24031 init: %24032) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24034 = stablehlo.multiply %23123, %23123 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24035 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24036 = stablehlo.reduce(%24034 init: %24035) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24037 = stablehlo.multiply %23126, %23126 : tensor<1x2048xf32> loc(#loc5242)
    %24038 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24039 = stablehlo.reduce(%24037 init: %24038) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24040 = stablehlo.multiply %23129, %23129 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24041 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24042 = stablehlo.reduce(%24040 init: %24041) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24043 = stablehlo.multiply %23132, %23132 : tensor<1x2048xf32> loc(#loc5242)
    %24044 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24045 = stablehlo.reduce(%24043 init: %24044) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24046 = stablehlo.multiply %23135, %23135 : tensor<1x2048xf32> loc(#loc5242)
    %24047 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24048 = stablehlo.reduce(%24046 init: %24047) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24049 = stablehlo.multiply %23138, %23138 : tensor<1x2048xf32> loc(#loc5242)
    %24050 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24051 = stablehlo.reduce(%24049 init: %24050) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24052 = stablehlo.multiply %23141, %23141 : tensor<1x2048xf32> loc(#loc5242)
    %24053 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24054 = stablehlo.reduce(%24052 init: %24053) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24055 = stablehlo.multiply %23144, %23144 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24056 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24057 = stablehlo.reduce(%24055 init: %24056) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24058 = stablehlo.multiply %23147, %23147 : tensor<1x64xf32> loc(#loc5238)
    %24059 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24060 = stablehlo.reduce(%24058 init: %24059) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24061 = stablehlo.multiply %23150, %23150 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24062 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24063 = stablehlo.reduce(%24061 init: %24062) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24064 = stablehlo.multiply %23153, %23153 : tensor<1x8192xf32> loc(#loc5238)
    %24065 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24066 = stablehlo.reduce(%24064 init: %24065) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24067 = stablehlo.multiply %23156, %23156 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24068 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24069 = stablehlo.reduce(%24067 init: %24068) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24070 = stablehlo.multiply %23159, %23159 : tensor<1x2048xf32> loc(#loc5242)
    %24071 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24072 = stablehlo.reduce(%24070 init: %24071) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24073 = stablehlo.multiply %23162, %23162 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24074 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24075 = stablehlo.reduce(%24073 init: %24074) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24076 = stablehlo.multiply %23165, %23165 : tensor<1x2048xf32> loc(#loc5242)
    %24077 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24078 = stablehlo.reduce(%24076 init: %24077) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24079 = stablehlo.multiply %23168, %23168 : tensor<1x2048xf32> loc(#loc5242)
    %24080 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24081 = stablehlo.reduce(%24079 init: %24080) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24082 = stablehlo.multiply %23171, %23171 : tensor<1x2048xf32> loc(#loc5242)
    %24083 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24084 = stablehlo.reduce(%24082 init: %24083) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24085 = stablehlo.multiply %23174, %23174 : tensor<1x2048xf32> loc(#loc5242)
    %24086 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24087 = stablehlo.reduce(%24085 init: %24086) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24088 = stablehlo.multiply %23177, %23177 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24089 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24090 = stablehlo.reduce(%24088 init: %24089) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24091 = stablehlo.multiply %23180, %23180 : tensor<1x64xf32> loc(#loc5238)
    %24092 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24093 = stablehlo.reduce(%24091 init: %24092) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24094 = stablehlo.multiply %23183, %23183 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24095 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24096 = stablehlo.reduce(%24094 init: %24095) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24097 = stablehlo.multiply %23186, %23186 : tensor<1x8192xf32> loc(#loc5238)
    %24098 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24099 = stablehlo.reduce(%24097 init: %24098) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24100 = stablehlo.multiply %23189, %23189 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24101 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24102 = stablehlo.reduce(%24100 init: %24101) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24103 = stablehlo.multiply %23192, %23192 : tensor<1x2048xf32> loc(#loc5242)
    %24104 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24105 = stablehlo.reduce(%24103 init: %24104) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24106 = stablehlo.multiply %23195, %23195 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24107 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24108 = stablehlo.reduce(%24106 init: %24107) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24109 = stablehlo.multiply %23198, %23198 : tensor<1x2048xf32> loc(#loc5242)
    %24110 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24111 = stablehlo.reduce(%24109 init: %24110) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24112 = stablehlo.multiply %23201, %23201 : tensor<1x2048xf32> loc(#loc5242)
    %24113 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24114 = stablehlo.reduce(%24112 init: %24113) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24115 = stablehlo.multiply %23204, %23204 : tensor<1x2048xf32> loc(#loc5242)
    %24116 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24117 = stablehlo.reduce(%24115 init: %24116) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24118 = stablehlo.multiply %23207, %23207 : tensor<1x2048xf32> loc(#loc5242)
    %24119 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24120 = stablehlo.reduce(%24118 init: %24119) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24121 = stablehlo.multiply %23210, %23210 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24122 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24123 = stablehlo.reduce(%24121 init: %24122) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24124 = stablehlo.multiply %23213, %23213 : tensor<1x64xf32> loc(#loc5238)
    %24125 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24126 = stablehlo.reduce(%24124 init: %24125) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24127 = stablehlo.multiply %23216, %23216 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24128 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24129 = stablehlo.reduce(%24127 init: %24128) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24130 = stablehlo.multiply %23219, %23219 : tensor<1x8192xf32> loc(#loc5238)
    %24131 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24132 = stablehlo.reduce(%24130 init: %24131) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24133 = stablehlo.multiply %23222, %23222 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24134 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24135 = stablehlo.reduce(%24133 init: %24134) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24136 = stablehlo.multiply %23225, %23225 : tensor<1x2048xf32> loc(#loc5242)
    %24137 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24138 = stablehlo.reduce(%24136 init: %24137) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24139 = stablehlo.multiply %23228, %23228 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24140 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24141 = stablehlo.reduce(%24139 init: %24140) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24142 = stablehlo.multiply %23231, %23231 : tensor<1x2048xf32> loc(#loc5242)
    %24143 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24144 = stablehlo.reduce(%24142 init: %24143) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24145 = stablehlo.multiply %23234, %23234 : tensor<1x2048xf32> loc(#loc5242)
    %24146 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24147 = stablehlo.reduce(%24145 init: %24146) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24148 = stablehlo.multiply %23237, %23237 : tensor<1x2048xf32> loc(#loc5242)
    %24149 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24150 = stablehlo.reduce(%24148 init: %24149) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24151 = stablehlo.multiply %23240, %23240 : tensor<1x2048xf32> loc(#loc5242)
    %24152 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24153 = stablehlo.reduce(%24151 init: %24152) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24154 = stablehlo.multiply %23243, %23243 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24155 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24156 = stablehlo.reduce(%24154 init: %24155) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24157 = stablehlo.multiply %23246, %23246 : tensor<1x64xf32> loc(#loc5238)
    %24158 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24159 = stablehlo.reduce(%24157 init: %24158) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24160 = stablehlo.multiply %23249, %23249 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24161 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24162 = stablehlo.reduce(%24160 init: %24161) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24163 = stablehlo.multiply %23252, %23252 : tensor<1x8192xf32> loc(#loc5238)
    %24164 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24165 = stablehlo.reduce(%24163 init: %24164) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24166 = stablehlo.multiply %23255, %23255 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24167 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24168 = stablehlo.reduce(%24166 init: %24167) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24169 = stablehlo.multiply %23258, %23258 : tensor<1x2048xf32> loc(#loc5242)
    %24170 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24171 = stablehlo.reduce(%24169 init: %24170) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24172 = stablehlo.multiply %23261, %23261 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24173 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24174 = stablehlo.reduce(%24172 init: %24173) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24175 = stablehlo.multiply %23264, %23264 : tensor<1x2048xf32> loc(#loc5242)
    %24176 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24177 = stablehlo.reduce(%24175 init: %24176) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24178 = stablehlo.multiply %23267, %23267 : tensor<1x2048xf32> loc(#loc5242)
    %24179 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24180 = stablehlo.reduce(%24178 init: %24179) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24181 = stablehlo.multiply %23270, %23270 : tensor<1x2048xf32> loc(#loc5242)
    %24182 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24183 = stablehlo.reduce(%24181 init: %24182) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24184 = stablehlo.multiply %23273, %23273 : tensor<1x2048xf32> loc(#loc5242)
    %24185 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24186 = stablehlo.reduce(%24184 init: %24185) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24187 = stablehlo.multiply %23276, %23276 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24188 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24189 = stablehlo.reduce(%24187 init: %24188) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24190 = stablehlo.multiply %23279, %23279 : tensor<1x64xf32> loc(#loc5238)
    %24191 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24192 = stablehlo.reduce(%24190 init: %24191) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24193 = stablehlo.multiply %23282, %23282 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24194 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24195 = stablehlo.reduce(%24193 init: %24194) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24196 = stablehlo.multiply %23285, %23285 : tensor<1x8192xf32> loc(#loc5238)
    %24197 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24198 = stablehlo.reduce(%24196 init: %24197) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24199 = stablehlo.multiply %23288, %23288 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24200 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24201 = stablehlo.reduce(%24199 init: %24200) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24202 = stablehlo.multiply %23291, %23291 : tensor<1x2048xf32> loc(#loc5242)
    %24203 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24204 = stablehlo.reduce(%24202 init: %24203) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24205 = stablehlo.multiply %23294, %23294 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24206 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24207 = stablehlo.reduce(%24205 init: %24206) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24208 = stablehlo.multiply %23297, %23297 : tensor<1x2048xf32> loc(#loc5242)
    %24209 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24210 = stablehlo.reduce(%24208 init: %24209) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24211 = stablehlo.multiply %23300, %23300 : tensor<1x2048xf32> loc(#loc5242)
    %24212 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24213 = stablehlo.reduce(%24211 init: %24212) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24214 = stablehlo.multiply %23303, %23303 : tensor<1x2048xf32> loc(#loc5242)
    %24215 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24216 = stablehlo.reduce(%24214 init: %24215) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24217 = stablehlo.multiply %23306, %23306 : tensor<1x2048xf32> loc(#loc5242)
    %24218 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24219 = stablehlo.reduce(%24217 init: %24218) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24220 = stablehlo.multiply %23309, %23309 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24221 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24222 = stablehlo.reduce(%24220 init: %24221) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24223 = stablehlo.multiply %23312, %23312 : tensor<1x64xf32> loc(#loc5238)
    %24224 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24225 = stablehlo.reduce(%24223 init: %24224) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24226 = stablehlo.multiply %23315, %23315 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24227 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24228 = stablehlo.reduce(%24226 init: %24227) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24229 = stablehlo.multiply %23318, %23318 : tensor<1x8192xf32> loc(#loc5238)
    %24230 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24231 = stablehlo.reduce(%24229 init: %24230) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24232 = stablehlo.multiply %23321, %23321 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24233 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24234 = stablehlo.reduce(%24232 init: %24233) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24235 = stablehlo.multiply %23324, %23324 : tensor<1x2048xf32> loc(#loc5242)
    %24236 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24237 = stablehlo.reduce(%24235 init: %24236) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24238 = stablehlo.multiply %23327, %23327 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24239 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24240 = stablehlo.reduce(%24238 init: %24239) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24241 = stablehlo.multiply %23330, %23330 : tensor<1x2048xf32> loc(#loc5242)
    %24242 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24243 = stablehlo.reduce(%24241 init: %24242) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24244 = stablehlo.multiply %23333, %23333 : tensor<1x2048xf32> loc(#loc5242)
    %24245 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24246 = stablehlo.reduce(%24244 init: %24245) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24247 = stablehlo.multiply %23336, %23336 : tensor<1x2048xf32> loc(#loc5242)
    %24248 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24249 = stablehlo.reduce(%24247 init: %24248) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24250 = stablehlo.multiply %23339, %23339 : tensor<1x2048xf32> loc(#loc5242)
    %24251 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24252 = stablehlo.reduce(%24250 init: %24251) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24253 = stablehlo.multiply %23342, %23342 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24254 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24255 = stablehlo.reduce(%24253 init: %24254) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24256 = stablehlo.multiply %23345, %23345 : tensor<1x64xf32> loc(#loc5238)
    %24257 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24258 = stablehlo.reduce(%24256 init: %24257) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24259 = stablehlo.multiply %23348, %23348 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24260 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24261 = stablehlo.reduce(%24259 init: %24260) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24262 = stablehlo.multiply %23351, %23351 : tensor<1x8192xf32> loc(#loc5238)
    %24263 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24264 = stablehlo.reduce(%24262 init: %24263) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24265 = stablehlo.multiply %23354, %23354 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24266 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24267 = stablehlo.reduce(%24265 init: %24266) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24268 = stablehlo.multiply %23357, %23357 : tensor<1x2048xf32> loc(#loc5242)
    %24269 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24270 = stablehlo.reduce(%24268 init: %24269) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24271 = stablehlo.multiply %23360, %23360 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24272 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24273 = stablehlo.reduce(%24271 init: %24272) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24274 = stablehlo.multiply %23363, %23363 : tensor<1x2048xf32> loc(#loc5242)
    %24275 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24276 = stablehlo.reduce(%24274 init: %24275) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24277 = stablehlo.multiply %23366, %23366 : tensor<1x2048xf32> loc(#loc5242)
    %24278 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24279 = stablehlo.reduce(%24277 init: %24278) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24280 = stablehlo.multiply %23369, %23369 : tensor<1x2048xf32> loc(#loc5242)
    %24281 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24282 = stablehlo.reduce(%24280 init: %24281) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24283 = stablehlo.multiply %23372, %23372 : tensor<1x2048xf32> loc(#loc5242)
    %24284 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24285 = stablehlo.reduce(%24283 init: %24284) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24286 = stablehlo.multiply %23375, %23375 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24287 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24288 = stablehlo.reduce(%24286 init: %24287) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24289 = stablehlo.multiply %23378, %23378 : tensor<1x64xf32> loc(#loc5238)
    %24290 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24291 = stablehlo.reduce(%24289 init: %24290) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24292 = stablehlo.multiply %23381, %23381 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24293 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24294 = stablehlo.reduce(%24292 init: %24293) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24295 = stablehlo.multiply %23384, %23384 : tensor<1x8192xf32> loc(#loc5238)
    %24296 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24297 = stablehlo.reduce(%24295 init: %24296) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24298 = stablehlo.multiply %23387, %23387 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24299 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24300 = stablehlo.reduce(%24298 init: %24299) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24301 = stablehlo.multiply %23390, %23390 : tensor<1x2048xf32> loc(#loc5242)
    %24302 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24303 = stablehlo.reduce(%24301 init: %24302) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24304 = stablehlo.multiply %23393, %23393 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24305 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24306 = stablehlo.reduce(%24304 init: %24305) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24307 = stablehlo.multiply %23396, %23396 : tensor<1x2048xf32> loc(#loc5242)
    %24308 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24309 = stablehlo.reduce(%24307 init: %24308) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24310 = stablehlo.multiply %23399, %23399 : tensor<1x2048xf32> loc(#loc5242)
    %24311 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24312 = stablehlo.reduce(%24310 init: %24311) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24313 = stablehlo.multiply %23402, %23402 : tensor<1x2048xf32> loc(#loc5242)
    %24314 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24315 = stablehlo.reduce(%24313 init: %24314) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24316 = stablehlo.multiply %23405, %23405 : tensor<1x2048xf32> loc(#loc5242)
    %24317 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24318 = stablehlo.reduce(%24316 init: %24317) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24319 = stablehlo.multiply %23408, %23408 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24320 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24321 = stablehlo.reduce(%24319 init: %24320) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24322 = stablehlo.multiply %23411, %23411 : tensor<1x64xf32> loc(#loc5238)
    %24323 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24324 = stablehlo.reduce(%24322 init: %24323) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24325 = stablehlo.multiply %23414, %23414 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24326 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24327 = stablehlo.reduce(%24325 init: %24326) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24328 = stablehlo.multiply %23417, %23417 : tensor<1x8192xf32> loc(#loc5238)
    %24329 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24330 = stablehlo.reduce(%24328 init: %24329) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24331 = stablehlo.multiply %23420, %23420 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24332 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24333 = stablehlo.reduce(%24331 init: %24332) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24334 = stablehlo.multiply %23423, %23423 : tensor<1x2048xf32> loc(#loc5242)
    %24335 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24336 = stablehlo.reduce(%24334 init: %24335) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24337 = stablehlo.multiply %23426, %23426 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24338 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24339 = stablehlo.reduce(%24337 init: %24338) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24340 = stablehlo.multiply %23429, %23429 : tensor<1x2048xf32> loc(#loc5242)
    %24341 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24342 = stablehlo.reduce(%24340 init: %24341) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24343 = stablehlo.multiply %23432, %23432 : tensor<1x2048xf32> loc(#loc5242)
    %24344 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24345 = stablehlo.reduce(%24343 init: %24344) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24346 = stablehlo.multiply %23435, %23435 : tensor<1x2048xf32> loc(#loc5242)
    %24347 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24348 = stablehlo.reduce(%24346 init: %24347) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24349 = stablehlo.multiply %23438, %23438 : tensor<1x2048xf32> loc(#loc5242)
    %24350 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24351 = stablehlo.reduce(%24349 init: %24350) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24352 = stablehlo.multiply %23441, %23441 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24353 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24354 = stablehlo.reduce(%24352 init: %24353) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24355 = stablehlo.multiply %23444, %23444 : tensor<1x64xf32> loc(#loc5238)
    %24356 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24357 = stablehlo.reduce(%24355 init: %24356) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24358 = stablehlo.multiply %23447, %23447 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24359 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24360 = stablehlo.reduce(%24358 init: %24359) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24361 = stablehlo.multiply %23450, %23450 : tensor<1x8192xf32> loc(#loc5238)
    %24362 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24363 = stablehlo.reduce(%24361 init: %24362) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24364 = stablehlo.multiply %23453, %23453 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24365 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24366 = stablehlo.reduce(%24364 init: %24365) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24367 = stablehlo.multiply %23456, %23456 : tensor<1x2048xf32> loc(#loc5242)
    %24368 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24369 = stablehlo.reduce(%24367 init: %24368) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24370 = stablehlo.multiply %23459, %23459 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24371 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24372 = stablehlo.reduce(%24370 init: %24371) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24373 = stablehlo.multiply %23462, %23462 : tensor<1x2048xf32> loc(#loc5242)
    %24374 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24375 = stablehlo.reduce(%24373 init: %24374) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24376 = stablehlo.multiply %23465, %23465 : tensor<1x2048xf32> loc(#loc5242)
    %24377 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24378 = stablehlo.reduce(%24376 init: %24377) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24379 = stablehlo.multiply %23468, %23468 : tensor<1x2048xf32> loc(#loc5242)
    %24380 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24381 = stablehlo.reduce(%24379 init: %24380) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24382 = stablehlo.multiply %23471, %23471 : tensor<1x2048xf32> loc(#loc5242)
    %24383 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24384 = stablehlo.reduce(%24382 init: %24383) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24385 = stablehlo.multiply %23474, %23474 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24386 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24387 = stablehlo.reduce(%24385 init: %24386) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24388 = stablehlo.multiply %23477, %23477 : tensor<1x64xf32> loc(#loc5238)
    %24389 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24390 = stablehlo.reduce(%24388 init: %24389) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24391 = stablehlo.multiply %23480, %23480 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24392 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24393 = stablehlo.reduce(%24391 init: %24392) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24394 = stablehlo.multiply %23483, %23483 : tensor<1x8192xf32> loc(#loc5238)
    %24395 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24396 = stablehlo.reduce(%24394 init: %24395) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24397 = stablehlo.multiply %23486, %23486 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24398 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24399 = stablehlo.reduce(%24397 init: %24398) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24400 = stablehlo.multiply %23489, %23489 : tensor<1x2048xf32> loc(#loc5242)
    %24401 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24402 = stablehlo.reduce(%24400 init: %24401) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24403 = stablehlo.multiply %23492, %23492 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24404 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24405 = stablehlo.reduce(%24403 init: %24404) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24406 = stablehlo.multiply %23495, %23495 : tensor<1x2048xf32> loc(#loc5242)
    %24407 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24408 = stablehlo.reduce(%24406 init: %24407) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24409 = stablehlo.multiply %23498, %23498 : tensor<1x2048xf32> loc(#loc5242)
    %24410 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24411 = stablehlo.reduce(%24409 init: %24410) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24412 = stablehlo.multiply %23501, %23501 : tensor<1x2048xf32> loc(#loc5242)
    %24413 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24414 = stablehlo.reduce(%24412 init: %24413) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24415 = stablehlo.multiply %23504, %23504 : tensor<1x2048xf32> loc(#loc5242)
    %24416 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24417 = stablehlo.reduce(%24415 init: %24416) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24418 = stablehlo.multiply %23507, %23507 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24419 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24420 = stablehlo.reduce(%24418 init: %24419) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24421 = stablehlo.multiply %23510, %23510 : tensor<1x64xf32> loc(#loc5238)
    %24422 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24423 = stablehlo.reduce(%24421 init: %24422) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24424 = stablehlo.multiply %23513, %23513 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24425 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24426 = stablehlo.reduce(%24424 init: %24425) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24427 = stablehlo.multiply %23516, %23516 : tensor<1x8192xf32> loc(#loc5238)
    %24428 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24429 = stablehlo.reduce(%24427 init: %24428) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24430 = stablehlo.multiply %23519, %23519 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24431 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24432 = stablehlo.reduce(%24430 init: %24431) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24433 = stablehlo.multiply %23522, %23522 : tensor<1x2048xf32> loc(#loc5242)
    %24434 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24435 = stablehlo.reduce(%24433 init: %24434) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24436 = stablehlo.multiply %23525, %23525 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24437 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24438 = stablehlo.reduce(%24436 init: %24437) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24439 = stablehlo.multiply %23528, %23528 : tensor<1x2048xf32> loc(#loc5242)
    %24440 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24441 = stablehlo.reduce(%24439 init: %24440) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24442 = stablehlo.multiply %23531, %23531 : tensor<1x2048xf32> loc(#loc5242)
    %24443 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24444 = stablehlo.reduce(%24442 init: %24443) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24445 = stablehlo.multiply %23534, %23534 : tensor<1x2048xf32> loc(#loc5242)
    %24446 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24447 = stablehlo.reduce(%24445 init: %24446) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24448 = stablehlo.multiply %23537, %23537 : tensor<1x2048xf32> loc(#loc5242)
    %24449 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24450 = stablehlo.reduce(%24448 init: %24449) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24451 = stablehlo.multiply %23540, %23540 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24452 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24453 = stablehlo.reduce(%24451 init: %24452) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24454 = stablehlo.multiply %23543, %23543 : tensor<1x64xf32> loc(#loc5238)
    %24455 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24456 = stablehlo.reduce(%24454 init: %24455) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24457 = stablehlo.multiply %23546, %23546 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24458 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24459 = stablehlo.reduce(%24457 init: %24458) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24460 = stablehlo.broadcast_in_dim %23919, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24461 = stablehlo.broadcast_in_dim %23922, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24462 = stablehlo.broadcast_in_dim %23925, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24463 = stablehlo.broadcast_in_dim %23928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24464 = stablehlo.broadcast_in_dim %23931, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24465 = stablehlo.broadcast_in_dim %23934, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24466 = stablehlo.broadcast_in_dim %23937, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24467 = stablehlo.broadcast_in_dim %23940, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24468 = stablehlo.broadcast_in_dim %23943, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24469 = stablehlo.broadcast_in_dim %23946, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24470 = stablehlo.broadcast_in_dim %23949, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24471 = stablehlo.broadcast_in_dim %23952, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24472 = stablehlo.broadcast_in_dim %23955, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24473 = stablehlo.broadcast_in_dim %23958, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24474 = stablehlo.broadcast_in_dim %23961, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24475 = stablehlo.broadcast_in_dim %23964, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24476 = stablehlo.broadcast_in_dim %23967, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24477 = stablehlo.broadcast_in_dim %23970, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24478 = stablehlo.broadcast_in_dim %23973, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24479 = stablehlo.broadcast_in_dim %23976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24480 = stablehlo.broadcast_in_dim %23979, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24481 = stablehlo.broadcast_in_dim %23982, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24482 = stablehlo.broadcast_in_dim %23985, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24483 = stablehlo.broadcast_in_dim %23988, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24484 = stablehlo.broadcast_in_dim %23991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24485 = stablehlo.broadcast_in_dim %23994, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24486 = stablehlo.broadcast_in_dim %23997, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24487 = stablehlo.broadcast_in_dim %24000, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24488 = stablehlo.broadcast_in_dim %24003, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24489 = stablehlo.broadcast_in_dim %24006, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24490 = stablehlo.broadcast_in_dim %24009, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24491 = stablehlo.broadcast_in_dim %24012, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24492 = stablehlo.broadcast_in_dim %24015, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24493 = stablehlo.broadcast_in_dim %24018, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24494 = stablehlo.broadcast_in_dim %24021, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24495 = stablehlo.broadcast_in_dim %24024, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24496 = stablehlo.broadcast_in_dim %24027, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24497 = stablehlo.broadcast_in_dim %24030, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24498 = stablehlo.broadcast_in_dim %24033, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24499 = stablehlo.broadcast_in_dim %24036, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24500 = stablehlo.broadcast_in_dim %24039, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24501 = stablehlo.broadcast_in_dim %24042, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24502 = stablehlo.broadcast_in_dim %24045, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24503 = stablehlo.broadcast_in_dim %24048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24504 = stablehlo.broadcast_in_dim %24051, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24505 = stablehlo.broadcast_in_dim %24054, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24506 = stablehlo.broadcast_in_dim %24057, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24507 = stablehlo.broadcast_in_dim %24060, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24508 = stablehlo.broadcast_in_dim %24063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24509 = stablehlo.broadcast_in_dim %24066, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24510 = stablehlo.broadcast_in_dim %24069, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24511 = stablehlo.broadcast_in_dim %24072, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24512 = stablehlo.broadcast_in_dim %24075, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24513 = stablehlo.broadcast_in_dim %24078, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24514 = stablehlo.broadcast_in_dim %24081, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24515 = stablehlo.broadcast_in_dim %24084, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24516 = stablehlo.broadcast_in_dim %24087, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24517 = stablehlo.broadcast_in_dim %24090, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24518 = stablehlo.broadcast_in_dim %24093, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24519 = stablehlo.broadcast_in_dim %24096, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24520 = stablehlo.broadcast_in_dim %24099, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24521 = stablehlo.broadcast_in_dim %24102, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24522 = stablehlo.broadcast_in_dim %24105, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24523 = stablehlo.broadcast_in_dim %24108, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24524 = stablehlo.broadcast_in_dim %24111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24525 = stablehlo.broadcast_in_dim %24114, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24526 = stablehlo.broadcast_in_dim %24117, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24527 = stablehlo.broadcast_in_dim %24120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24528 = stablehlo.broadcast_in_dim %24123, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24529 = stablehlo.broadcast_in_dim %24126, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24530 = stablehlo.broadcast_in_dim %24129, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24531 = stablehlo.broadcast_in_dim %24132, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24532 = stablehlo.broadcast_in_dim %24135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24533 = stablehlo.broadcast_in_dim %24138, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24534 = stablehlo.broadcast_in_dim %24141, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24535 = stablehlo.broadcast_in_dim %24144, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24536 = stablehlo.broadcast_in_dim %24147, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24537 = stablehlo.broadcast_in_dim %24150, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24538 = stablehlo.broadcast_in_dim %24153, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24539 = stablehlo.broadcast_in_dim %24156, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24540 = stablehlo.broadcast_in_dim %24159, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24541 = stablehlo.broadcast_in_dim %24162, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24542 = stablehlo.broadcast_in_dim %24165, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24543 = stablehlo.broadcast_in_dim %24168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24544 = stablehlo.broadcast_in_dim %24171, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24545 = stablehlo.broadcast_in_dim %24174, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24546 = stablehlo.broadcast_in_dim %24177, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24547 = stablehlo.broadcast_in_dim %24180, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24548 = stablehlo.broadcast_in_dim %24183, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24549 = stablehlo.broadcast_in_dim %24186, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24550 = stablehlo.broadcast_in_dim %24189, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24551 = stablehlo.broadcast_in_dim %24192, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24552 = stablehlo.broadcast_in_dim %24195, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24553 = stablehlo.broadcast_in_dim %24198, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24554 = stablehlo.broadcast_in_dim %24201, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24555 = stablehlo.broadcast_in_dim %24204, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24556 = stablehlo.broadcast_in_dim %24207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24557 = stablehlo.broadcast_in_dim %24210, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24558 = stablehlo.broadcast_in_dim %24213, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24559 = stablehlo.broadcast_in_dim %24216, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24560 = stablehlo.broadcast_in_dim %24219, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24561 = stablehlo.broadcast_in_dim %24222, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24562 = stablehlo.broadcast_in_dim %24225, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24563 = stablehlo.broadcast_in_dim %24228, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24564 = stablehlo.broadcast_in_dim %24231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24565 = stablehlo.broadcast_in_dim %24234, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24566 = stablehlo.broadcast_in_dim %24237, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24567 = stablehlo.broadcast_in_dim %24240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24568 = stablehlo.broadcast_in_dim %24243, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24569 = stablehlo.broadcast_in_dim %24246, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24570 = stablehlo.broadcast_in_dim %24249, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24571 = stablehlo.broadcast_in_dim %24252, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24572 = stablehlo.broadcast_in_dim %24255, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24573 = stablehlo.broadcast_in_dim %24258, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24574 = stablehlo.broadcast_in_dim %24261, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24575 = stablehlo.broadcast_in_dim %24264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24576 = stablehlo.broadcast_in_dim %24267, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24577 = stablehlo.broadcast_in_dim %24270, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24578 = stablehlo.broadcast_in_dim %24273, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24579 = stablehlo.broadcast_in_dim %24276, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24580 = stablehlo.broadcast_in_dim %24279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24581 = stablehlo.broadcast_in_dim %24282, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24582 = stablehlo.broadcast_in_dim %24285, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24583 = stablehlo.broadcast_in_dim %24288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24584 = stablehlo.broadcast_in_dim %24291, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24585 = stablehlo.broadcast_in_dim %24294, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24586 = stablehlo.broadcast_in_dim %24297, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24587 = stablehlo.broadcast_in_dim %24300, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24588 = stablehlo.broadcast_in_dim %24303, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24589 = stablehlo.broadcast_in_dim %24306, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24590 = stablehlo.broadcast_in_dim %24309, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24591 = stablehlo.broadcast_in_dim %24312, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24592 = stablehlo.broadcast_in_dim %24315, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24593 = stablehlo.broadcast_in_dim %24318, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24594 = stablehlo.broadcast_in_dim %24321, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24595 = stablehlo.broadcast_in_dim %24324, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24596 = stablehlo.broadcast_in_dim %24327, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24597 = stablehlo.broadcast_in_dim %24330, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24598 = stablehlo.broadcast_in_dim %24333, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24599 = stablehlo.broadcast_in_dim %24336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24600 = stablehlo.broadcast_in_dim %24339, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24601 = stablehlo.broadcast_in_dim %24342, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24602 = stablehlo.broadcast_in_dim %24345, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24603 = stablehlo.broadcast_in_dim %24348, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24604 = stablehlo.broadcast_in_dim %24351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24605 = stablehlo.broadcast_in_dim %24354, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24606 = stablehlo.broadcast_in_dim %24357, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24607 = stablehlo.broadcast_in_dim %24360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24608 = stablehlo.broadcast_in_dim %24363, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24609 = stablehlo.broadcast_in_dim %24366, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24610 = stablehlo.broadcast_in_dim %24369, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24611 = stablehlo.broadcast_in_dim %24372, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24612 = stablehlo.broadcast_in_dim %24375, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24613 = stablehlo.broadcast_in_dim %24378, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24614 = stablehlo.broadcast_in_dim %24381, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24615 = stablehlo.broadcast_in_dim %24384, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24616 = stablehlo.broadcast_in_dim %24387, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24617 = stablehlo.broadcast_in_dim %24390, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24618 = stablehlo.broadcast_in_dim %24393, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24619 = stablehlo.broadcast_in_dim %24396, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24620 = stablehlo.broadcast_in_dim %24399, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24621 = stablehlo.broadcast_in_dim %24402, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24622 = stablehlo.broadcast_in_dim %24405, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24623 = stablehlo.broadcast_in_dim %24408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24624 = stablehlo.broadcast_in_dim %24411, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24625 = stablehlo.broadcast_in_dim %24414, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24626 = stablehlo.broadcast_in_dim %24417, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24627 = stablehlo.broadcast_in_dim %24420, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24628 = stablehlo.broadcast_in_dim %24423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24629 = stablehlo.broadcast_in_dim %24426, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24630 = stablehlo.broadcast_in_dim %24429, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24631 = stablehlo.broadcast_in_dim %24432, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24632 = stablehlo.broadcast_in_dim %24435, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24633 = stablehlo.broadcast_in_dim %24438, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24634 = stablehlo.broadcast_in_dim %24441, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24635 = stablehlo.broadcast_in_dim %24444, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24636 = stablehlo.broadcast_in_dim %24447, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24637 = stablehlo.broadcast_in_dim %24450, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24638 = stablehlo.broadcast_in_dim %24453, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24639 = stablehlo.broadcast_in_dim %24456, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24640 = stablehlo.broadcast_in_dim %24459, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5246)
    %24641 = stablehlo.concatenate %24460, %24461, %24462, %24463, %24464, %24465, %24466, %24467, %24468, %24469, %24470, %24471, %24472, %24473, %24474, %24475, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24642 = stablehlo.concatenate %24476, %24477, %24478, %24479, %24480, %24481, %24482, %24483, %24484, %24485, %24486, %24487, %24488, %24489, %24490, %24491, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24643 = stablehlo.concatenate %24492, %24493, %24494, %24495, %24496, %24497, %24498, %24499, %24500, %24501, %24502, %24503, %24504, %24505, %24506, %24507, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24644 = stablehlo.concatenate %24508, %24509, %24510, %24511, %24512, %24513, %24514, %24515, %24516, %24517, %24518, %24519, %24520, %24521, %24522, %24523, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24645 = stablehlo.concatenate %24524, %24525, %24526, %24527, %24528, %24529, %24530, %24531, %24532, %24533, %24534, %24535, %24536, %24537, %24538, %24539, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24646 = stablehlo.concatenate %24540, %24541, %24542, %24543, %24544, %24545, %24546, %24547, %24548, %24549, %24550, %24551, %24552, %24553, %24554, %24555, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24647 = stablehlo.concatenate %24556, %24557, %24558, %24559, %24560, %24561, %24562, %24563, %24564, %24565, %24566, %24567, %24568, %24569, %24570, %24571, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24648 = stablehlo.concatenate %24572, %24573, %24574, %24575, %24576, %24577, %24578, %24579, %24580, %24581, %24582, %24583, %24584, %24585, %24586, %24587, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24649 = stablehlo.concatenate %24588, %24589, %24590, %24591, %24592, %24593, %24594, %24595, %24596, %24597, %24598, %24599, %24600, %24601, %24602, %24603, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24650 = stablehlo.concatenate %24604, %24605, %24606, %24607, %24608, %24609, %24610, %24611, %24612, %24613, %24614, %24615, %24616, %24617, %24618, %24619, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24651 = stablehlo.concatenate %24620, %24621, %24622, %24623, %24624, %24625, %24626, %24627, %24628, %24629, %24630, %24631, %24632, %24633, %24634, %24635, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5247)
    %24652 = stablehlo.concatenate %24636, %24637, %24638, %24639, %24640, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5xf32> loc(#loc5247)
    %24653 = stablehlo.concatenate %24641, %24642, %24643, %24644, %24645, %24646, %24647, %24648, %24649, %24650, %24651, %24652, dim = 0 : (tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<5xf32>) -> tensor<181xf32> loc(#loc5247)
    %24654 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5248)
    %24655 = stablehlo.reduce(%24653 init: %24654) across dimensions = [0] : (tensor<181xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5248)
      stablehlo.return %29929 : tensor<f32> loc(#loc5248)
    } loc(#loc5248)
    %24656 = stablehlo.sqrt %24655 : tensor<f32> loc(#loc5249)
    %24657 = stablehlo.multiply %arg1, %arg1 : tensor<2048xf32> loc(#loc5238)
    %24658 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %24659 = stablehlo.reduce(%24657 init: %24658) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %24660 = stablehlo.multiply %arg2, %arg2 : tensor<2048xf32> loc(#loc5238)
    %24661 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %24662 = stablehlo.reduce(%24660 init: %24661) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %24663 = stablehlo.multiply %arg3, %arg3 : tensor<2048x2048xf32> loc(#loc5238)
    %24664 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24665 = stablehlo.reduce(%24663 init: %24664) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24666 = stablehlo.multiply %arg4, %arg4 : tensor<51200xf32> loc(#loc5238)
    %24667 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %24668 = stablehlo.reduce(%24666 init: %24667) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %24669 = stablehlo.multiply %arg5, %arg5 : tensor<2048x51200xf32> loc(#loc5238)
    %24670 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24671 = stablehlo.reduce(%24669 init: %24670) across dimensions = [0, 1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24672 = stablehlo.multiply %arg6, %arg6 : tensor<1x8192xf32> loc(#loc5238)
    %24673 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24674 = stablehlo.reduce(%24672 init: %24673) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24675 = stablehlo.multiply %arg7, %arg7 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24676 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24677 = stablehlo.reduce(%24675 init: %24676) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24678 = stablehlo.multiply %arg8, %arg8 : tensor<1x2048xf32> loc(#loc5242)
    %24679 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24680 = stablehlo.reduce(%24678 init: %24679) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24681 = stablehlo.multiply %arg9, %arg9 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24682 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24683 = stablehlo.reduce(%24681 init: %24682) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24684 = stablehlo.multiply %arg10, %arg10 : tensor<1x2048xf32> loc(#loc5242)
    %24685 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24686 = stablehlo.reduce(%24684 init: %24685) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24687 = stablehlo.multiply %arg11, %arg11 : tensor<1x2048xf32> loc(#loc5242)
    %24688 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24689 = stablehlo.reduce(%24687 init: %24688) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24690 = stablehlo.multiply %arg12, %arg12 : tensor<1x2048xf32> loc(#loc5242)
    %24691 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24692 = stablehlo.reduce(%24690 init: %24691) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24693 = stablehlo.multiply %arg13, %arg13 : tensor<1x2048xf32> loc(#loc5242)
    %24694 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24695 = stablehlo.reduce(%24693 init: %24694) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24696 = stablehlo.multiply %arg14, %arg14 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24697 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24698 = stablehlo.reduce(%24696 init: %24697) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24699 = stablehlo.multiply %arg15, %arg15 : tensor<1x64xf32> loc(#loc5238)
    %24700 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24701 = stablehlo.reduce(%24699 init: %24700) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24702 = stablehlo.multiply %arg16, %arg16 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24703 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24704 = stablehlo.reduce(%24702 init: %24703) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24705 = stablehlo.multiply %arg17, %arg17 : tensor<1x8192xf32> loc(#loc5238)
    %24706 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24707 = stablehlo.reduce(%24705 init: %24706) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24708 = stablehlo.multiply %arg18, %arg18 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24709 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24710 = stablehlo.reduce(%24708 init: %24709) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24711 = stablehlo.multiply %arg19, %arg19 : tensor<1x2048xf32> loc(#loc5242)
    %24712 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24713 = stablehlo.reduce(%24711 init: %24712) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24714 = stablehlo.multiply %arg20, %arg20 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24715 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24716 = stablehlo.reduce(%24714 init: %24715) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24717 = stablehlo.multiply %arg21, %arg21 : tensor<1x2048xf32> loc(#loc5242)
    %24718 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24719 = stablehlo.reduce(%24717 init: %24718) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24720 = stablehlo.multiply %arg22, %arg22 : tensor<1x2048xf32> loc(#loc5242)
    %24721 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24722 = stablehlo.reduce(%24720 init: %24721) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24723 = stablehlo.multiply %arg23, %arg23 : tensor<1x2048xf32> loc(#loc5242)
    %24724 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24725 = stablehlo.reduce(%24723 init: %24724) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24726 = stablehlo.multiply %arg24, %arg24 : tensor<1x2048xf32> loc(#loc5242)
    %24727 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24728 = stablehlo.reduce(%24726 init: %24727) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24729 = stablehlo.multiply %arg25, %arg25 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24730 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24731 = stablehlo.reduce(%24729 init: %24730) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24732 = stablehlo.multiply %arg26, %arg26 : tensor<1x64xf32> loc(#loc5238)
    %24733 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24734 = stablehlo.reduce(%24732 init: %24733) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24735 = stablehlo.multiply %arg27, %arg27 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24736 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24737 = stablehlo.reduce(%24735 init: %24736) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24738 = stablehlo.multiply %arg28, %arg28 : tensor<1x8192xf32> loc(#loc5238)
    %24739 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24740 = stablehlo.reduce(%24738 init: %24739) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24741 = stablehlo.multiply %arg29, %arg29 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24742 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24743 = stablehlo.reduce(%24741 init: %24742) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24744 = stablehlo.multiply %arg30, %arg30 : tensor<1x2048xf32> loc(#loc5242)
    %24745 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24746 = stablehlo.reduce(%24744 init: %24745) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24747 = stablehlo.multiply %arg31, %arg31 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24748 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24749 = stablehlo.reduce(%24747 init: %24748) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24750 = stablehlo.multiply %arg32, %arg32 : tensor<1x2048xf32> loc(#loc5242)
    %24751 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24752 = stablehlo.reduce(%24750 init: %24751) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24753 = stablehlo.multiply %arg33, %arg33 : tensor<1x2048xf32> loc(#loc5242)
    %24754 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24755 = stablehlo.reduce(%24753 init: %24754) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24756 = stablehlo.multiply %arg34, %arg34 : tensor<1x2048xf32> loc(#loc5242)
    %24757 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24758 = stablehlo.reduce(%24756 init: %24757) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24759 = stablehlo.multiply %arg35, %arg35 : tensor<1x2048xf32> loc(#loc5242)
    %24760 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24761 = stablehlo.reduce(%24759 init: %24760) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24762 = stablehlo.multiply %arg36, %arg36 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24763 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24764 = stablehlo.reduce(%24762 init: %24763) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24765 = stablehlo.multiply %arg37, %arg37 : tensor<1x64xf32> loc(#loc5238)
    %24766 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24767 = stablehlo.reduce(%24765 init: %24766) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24768 = stablehlo.multiply %arg38, %arg38 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24769 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24770 = stablehlo.reduce(%24768 init: %24769) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24771 = stablehlo.multiply %arg39, %arg39 : tensor<1x8192xf32> loc(#loc5238)
    %24772 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24773 = stablehlo.reduce(%24771 init: %24772) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24774 = stablehlo.multiply %arg40, %arg40 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24775 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24776 = stablehlo.reduce(%24774 init: %24775) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24777 = stablehlo.multiply %arg41, %arg41 : tensor<1x2048xf32> loc(#loc5242)
    %24778 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24779 = stablehlo.reduce(%24777 init: %24778) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24780 = stablehlo.multiply %arg42, %arg42 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24781 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24782 = stablehlo.reduce(%24780 init: %24781) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24783 = stablehlo.multiply %arg43, %arg43 : tensor<1x2048xf32> loc(#loc5242)
    %24784 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24785 = stablehlo.reduce(%24783 init: %24784) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24786 = stablehlo.multiply %arg44, %arg44 : tensor<1x2048xf32> loc(#loc5242)
    %24787 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24788 = stablehlo.reduce(%24786 init: %24787) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24789 = stablehlo.multiply %arg45, %arg45 : tensor<1x2048xf32> loc(#loc5242)
    %24790 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24791 = stablehlo.reduce(%24789 init: %24790) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24792 = stablehlo.multiply %arg46, %arg46 : tensor<1x2048xf32> loc(#loc5242)
    %24793 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24794 = stablehlo.reduce(%24792 init: %24793) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24795 = stablehlo.multiply %arg47, %arg47 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24796 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24797 = stablehlo.reduce(%24795 init: %24796) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24798 = stablehlo.multiply %arg48, %arg48 : tensor<1x64xf32> loc(#loc5238)
    %24799 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24800 = stablehlo.reduce(%24798 init: %24799) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24801 = stablehlo.multiply %arg49, %arg49 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24802 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24803 = stablehlo.reduce(%24801 init: %24802) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24804 = stablehlo.multiply %arg50, %arg50 : tensor<1x8192xf32> loc(#loc5238)
    %24805 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24806 = stablehlo.reduce(%24804 init: %24805) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24807 = stablehlo.multiply %arg51, %arg51 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24808 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24809 = stablehlo.reduce(%24807 init: %24808) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24810 = stablehlo.multiply %arg52, %arg52 : tensor<1x2048xf32> loc(#loc5242)
    %24811 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24812 = stablehlo.reduce(%24810 init: %24811) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24813 = stablehlo.multiply %arg53, %arg53 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24814 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24815 = stablehlo.reduce(%24813 init: %24814) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24816 = stablehlo.multiply %arg54, %arg54 : tensor<1x2048xf32> loc(#loc5242)
    %24817 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24818 = stablehlo.reduce(%24816 init: %24817) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24819 = stablehlo.multiply %arg55, %arg55 : tensor<1x2048xf32> loc(#loc5242)
    %24820 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24821 = stablehlo.reduce(%24819 init: %24820) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24822 = stablehlo.multiply %arg56, %arg56 : tensor<1x2048xf32> loc(#loc5242)
    %24823 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24824 = stablehlo.reduce(%24822 init: %24823) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24825 = stablehlo.multiply %arg57, %arg57 : tensor<1x2048xf32> loc(#loc5242)
    %24826 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24827 = stablehlo.reduce(%24825 init: %24826) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24828 = stablehlo.multiply %arg58, %arg58 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24829 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24830 = stablehlo.reduce(%24828 init: %24829) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24831 = stablehlo.multiply %arg59, %arg59 : tensor<1x64xf32> loc(#loc5238)
    %24832 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24833 = stablehlo.reduce(%24831 init: %24832) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24834 = stablehlo.multiply %arg60, %arg60 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24835 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24836 = stablehlo.reduce(%24834 init: %24835) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24837 = stablehlo.multiply %arg61, %arg61 : tensor<1x8192xf32> loc(#loc5238)
    %24838 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24839 = stablehlo.reduce(%24837 init: %24838) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24840 = stablehlo.multiply %arg62, %arg62 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24841 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24842 = stablehlo.reduce(%24840 init: %24841) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24843 = stablehlo.multiply %arg63, %arg63 : tensor<1x2048xf32> loc(#loc5242)
    %24844 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24845 = stablehlo.reduce(%24843 init: %24844) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24846 = stablehlo.multiply %arg64, %arg64 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24847 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24848 = stablehlo.reduce(%24846 init: %24847) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24849 = stablehlo.multiply %arg65, %arg65 : tensor<1x2048xf32> loc(#loc5242)
    %24850 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24851 = stablehlo.reduce(%24849 init: %24850) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24852 = stablehlo.multiply %arg66, %arg66 : tensor<1x2048xf32> loc(#loc5242)
    %24853 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24854 = stablehlo.reduce(%24852 init: %24853) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24855 = stablehlo.multiply %arg67, %arg67 : tensor<1x2048xf32> loc(#loc5242)
    %24856 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24857 = stablehlo.reduce(%24855 init: %24856) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24858 = stablehlo.multiply %arg68, %arg68 : tensor<1x2048xf32> loc(#loc5242)
    %24859 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24860 = stablehlo.reduce(%24858 init: %24859) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24861 = stablehlo.multiply %arg69, %arg69 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24862 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24863 = stablehlo.reduce(%24861 init: %24862) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24864 = stablehlo.multiply %arg70, %arg70 : tensor<1x64xf32> loc(#loc5238)
    %24865 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24866 = stablehlo.reduce(%24864 init: %24865) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24867 = stablehlo.multiply %arg71, %arg71 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24868 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24869 = stablehlo.reduce(%24867 init: %24868) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24870 = stablehlo.multiply %arg72, %arg72 : tensor<1x8192xf32> loc(#loc5238)
    %24871 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24872 = stablehlo.reduce(%24870 init: %24871) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24873 = stablehlo.multiply %arg73, %arg73 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24874 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24875 = stablehlo.reduce(%24873 init: %24874) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24876 = stablehlo.multiply %arg74, %arg74 : tensor<1x2048xf32> loc(#loc5242)
    %24877 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24878 = stablehlo.reduce(%24876 init: %24877) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24879 = stablehlo.multiply %arg75, %arg75 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24880 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24881 = stablehlo.reduce(%24879 init: %24880) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24882 = stablehlo.multiply %arg76, %arg76 : tensor<1x2048xf32> loc(#loc5242)
    %24883 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24884 = stablehlo.reduce(%24882 init: %24883) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24885 = stablehlo.multiply %arg77, %arg77 : tensor<1x2048xf32> loc(#loc5242)
    %24886 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24887 = stablehlo.reduce(%24885 init: %24886) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24888 = stablehlo.multiply %arg78, %arg78 : tensor<1x2048xf32> loc(#loc5242)
    %24889 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24890 = stablehlo.reduce(%24888 init: %24889) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24891 = stablehlo.multiply %arg79, %arg79 : tensor<1x2048xf32> loc(#loc5242)
    %24892 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24893 = stablehlo.reduce(%24891 init: %24892) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24894 = stablehlo.multiply %arg80, %arg80 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24895 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24896 = stablehlo.reduce(%24894 init: %24895) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24897 = stablehlo.multiply %arg81, %arg81 : tensor<1x64xf32> loc(#loc5238)
    %24898 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24899 = stablehlo.reduce(%24897 init: %24898) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24900 = stablehlo.multiply %arg82, %arg82 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24901 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24902 = stablehlo.reduce(%24900 init: %24901) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24903 = stablehlo.multiply %arg83, %arg83 : tensor<1x8192xf32> loc(#loc5238)
    %24904 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24905 = stablehlo.reduce(%24903 init: %24904) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24906 = stablehlo.multiply %arg84, %arg84 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24907 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24908 = stablehlo.reduce(%24906 init: %24907) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24909 = stablehlo.multiply %arg85, %arg85 : tensor<1x2048xf32> loc(#loc5242)
    %24910 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24911 = stablehlo.reduce(%24909 init: %24910) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24912 = stablehlo.multiply %arg86, %arg86 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24913 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24914 = stablehlo.reduce(%24912 init: %24913) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24915 = stablehlo.multiply %arg87, %arg87 : tensor<1x2048xf32> loc(#loc5242)
    %24916 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24917 = stablehlo.reduce(%24915 init: %24916) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24918 = stablehlo.multiply %arg88, %arg88 : tensor<1x2048xf32> loc(#loc5242)
    %24919 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24920 = stablehlo.reduce(%24918 init: %24919) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24921 = stablehlo.multiply %arg89, %arg89 : tensor<1x2048xf32> loc(#loc5242)
    %24922 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24923 = stablehlo.reduce(%24921 init: %24922) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24924 = stablehlo.multiply %arg90, %arg90 : tensor<1x2048xf32> loc(#loc5242)
    %24925 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24926 = stablehlo.reduce(%24924 init: %24925) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24927 = stablehlo.multiply %arg91, %arg91 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24928 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24929 = stablehlo.reduce(%24927 init: %24928) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24930 = stablehlo.multiply %arg92, %arg92 : tensor<1x64xf32> loc(#loc5238)
    %24931 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24932 = stablehlo.reduce(%24930 init: %24931) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24933 = stablehlo.multiply %arg93, %arg93 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24934 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24935 = stablehlo.reduce(%24933 init: %24934) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24936 = stablehlo.multiply %arg94, %arg94 : tensor<1x8192xf32> loc(#loc5238)
    %24937 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24938 = stablehlo.reduce(%24936 init: %24937) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24939 = stablehlo.multiply %arg95, %arg95 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24940 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24941 = stablehlo.reduce(%24939 init: %24940) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24942 = stablehlo.multiply %arg96, %arg96 : tensor<1x2048xf32> loc(#loc5242)
    %24943 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24944 = stablehlo.reduce(%24942 init: %24943) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24945 = stablehlo.multiply %arg97, %arg97 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24946 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24947 = stablehlo.reduce(%24945 init: %24946) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24948 = stablehlo.multiply %arg98, %arg98 : tensor<1x2048xf32> loc(#loc5242)
    %24949 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24950 = stablehlo.reduce(%24948 init: %24949) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24951 = stablehlo.multiply %arg99, %arg99 : tensor<1x2048xf32> loc(#loc5242)
    %24952 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24953 = stablehlo.reduce(%24951 init: %24952) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24954 = stablehlo.multiply %arg100, %arg100 : tensor<1x2048xf32> loc(#loc5242)
    %24955 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24956 = stablehlo.reduce(%24954 init: %24955) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24957 = stablehlo.multiply %arg101, %arg101 : tensor<1x2048xf32> loc(#loc5242)
    %24958 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24959 = stablehlo.reduce(%24957 init: %24958) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24960 = stablehlo.multiply %arg102, %arg102 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24961 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24962 = stablehlo.reduce(%24960 init: %24961) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24963 = stablehlo.multiply %arg103, %arg103 : tensor<1x64xf32> loc(#loc5238)
    %24964 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24965 = stablehlo.reduce(%24963 init: %24964) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24966 = stablehlo.multiply %arg104, %arg104 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %24967 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %24968 = stablehlo.reduce(%24966 init: %24967) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %24969 = stablehlo.multiply %arg105, %arg105 : tensor<1x8192xf32> loc(#loc5238)
    %24970 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24971 = stablehlo.reduce(%24969 init: %24970) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24972 = stablehlo.multiply %arg106, %arg106 : tensor<1x2048x8192xf32> loc(#loc5238)
    %24973 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24974 = stablehlo.reduce(%24972 init: %24973) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24975 = stablehlo.multiply %arg107, %arg107 : tensor<1x2048xf32> loc(#loc5242)
    %24976 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24977 = stablehlo.reduce(%24975 init: %24976) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24978 = stablehlo.multiply %arg108, %arg108 : tensor<1x8192x2048xf32> loc(#loc5238)
    %24979 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %24980 = stablehlo.reduce(%24978 init: %24979) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %24981 = stablehlo.multiply %arg109, %arg109 : tensor<1x2048xf32> loc(#loc5242)
    %24982 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24983 = stablehlo.reduce(%24981 init: %24982) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24984 = stablehlo.multiply %arg110, %arg110 : tensor<1x2048xf32> loc(#loc5242)
    %24985 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24986 = stablehlo.reduce(%24984 init: %24985) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24987 = stablehlo.multiply %arg111, %arg111 : tensor<1x2048xf32> loc(#loc5242)
    %24988 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24989 = stablehlo.reduce(%24987 init: %24988) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24990 = stablehlo.multiply %arg112, %arg112 : tensor<1x2048xf32> loc(#loc5242)
    %24991 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %24992 = stablehlo.reduce(%24990 init: %24991) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %24993 = stablehlo.multiply %arg113, %arg113 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %24994 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %24995 = stablehlo.reduce(%24993 init: %24994) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %24996 = stablehlo.multiply %arg114, %arg114 : tensor<1x64xf32> loc(#loc5238)
    %24997 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %24998 = stablehlo.reduce(%24996 init: %24997) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %24999 = stablehlo.multiply %arg115, %arg115 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25000 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25001 = stablehlo.reduce(%24999 init: %25000) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25002 = stablehlo.multiply %arg116, %arg116 : tensor<1x8192xf32> loc(#loc5238)
    %25003 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25004 = stablehlo.reduce(%25002 init: %25003) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25005 = stablehlo.multiply %arg117, %arg117 : tensor<1x2048x8192xf32> loc(#loc5238)
    %25006 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25007 = stablehlo.reduce(%25005 init: %25006) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25008 = stablehlo.multiply %arg118, %arg118 : tensor<1x2048xf32> loc(#loc5242)
    %25009 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25010 = stablehlo.reduce(%25008 init: %25009) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25011 = stablehlo.multiply %arg119, %arg119 : tensor<1x8192x2048xf32> loc(#loc5238)
    %25012 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25013 = stablehlo.reduce(%25011 init: %25012) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25014 = stablehlo.multiply %arg120, %arg120 : tensor<1x2048xf32> loc(#loc5242)
    %25015 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25016 = stablehlo.reduce(%25014 init: %25015) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25017 = stablehlo.multiply %arg121, %arg121 : tensor<1x2048xf32> loc(#loc5242)
    %25018 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25019 = stablehlo.reduce(%25017 init: %25018) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25020 = stablehlo.multiply %arg122, %arg122 : tensor<1x2048xf32> loc(#loc5242)
    %25021 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25022 = stablehlo.reduce(%25020 init: %25021) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25023 = stablehlo.multiply %arg123, %arg123 : tensor<1x2048xf32> loc(#loc5242)
    %25024 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25025 = stablehlo.reduce(%25023 init: %25024) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25026 = stablehlo.multiply %arg124, %arg124 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %25027 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25028 = stablehlo.reduce(%25026 init: %25027) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25029 = stablehlo.multiply %arg125, %arg125 : tensor<1x64xf32> loc(#loc5238)
    %25030 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25031 = stablehlo.reduce(%25029 init: %25030) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25032 = stablehlo.multiply %arg126, %arg126 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25033 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25034 = stablehlo.reduce(%25032 init: %25033) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25035 = stablehlo.multiply %arg127, %arg127 : tensor<1x8192xf32> loc(#loc5238)
    %25036 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25037 = stablehlo.reduce(%25035 init: %25036) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25038 = stablehlo.multiply %arg128, %arg128 : tensor<1x2048x8192xf32> loc(#loc5238)
    %25039 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25040 = stablehlo.reduce(%25038 init: %25039) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25041 = stablehlo.multiply %arg129, %arg129 : tensor<1x2048xf32> loc(#loc5242)
    %25042 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25043 = stablehlo.reduce(%25041 init: %25042) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25044 = stablehlo.multiply %arg130, %arg130 : tensor<1x8192x2048xf32> loc(#loc5238)
    %25045 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25046 = stablehlo.reduce(%25044 init: %25045) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25047 = stablehlo.multiply %arg131, %arg131 : tensor<1x2048xf32> loc(#loc5242)
    %25048 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25049 = stablehlo.reduce(%25047 init: %25048) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25050 = stablehlo.multiply %arg132, %arg132 : tensor<1x2048xf32> loc(#loc5242)
    %25051 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25052 = stablehlo.reduce(%25050 init: %25051) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25053 = stablehlo.multiply %arg133, %arg133 : tensor<1x2048xf32> loc(#loc5242)
    %25054 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25055 = stablehlo.reduce(%25053 init: %25054) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25056 = stablehlo.multiply %arg134, %arg134 : tensor<1x2048xf32> loc(#loc5242)
    %25057 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25058 = stablehlo.reduce(%25056 init: %25057) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25059 = stablehlo.multiply %arg135, %arg135 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %25060 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25061 = stablehlo.reduce(%25059 init: %25060) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25062 = stablehlo.multiply %arg136, %arg136 : tensor<1x64xf32> loc(#loc5238)
    %25063 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25064 = stablehlo.reduce(%25062 init: %25063) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25065 = stablehlo.multiply %arg137, %arg137 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25066 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25067 = stablehlo.reduce(%25065 init: %25066) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25068 = stablehlo.multiply %arg138, %arg138 : tensor<1x8192xf32> loc(#loc5238)
    %25069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25070 = stablehlo.reduce(%25068 init: %25069) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25071 = stablehlo.multiply %arg139, %arg139 : tensor<1x2048x8192xf32> loc(#loc5238)
    %25072 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25073 = stablehlo.reduce(%25071 init: %25072) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25074 = stablehlo.multiply %arg140, %arg140 : tensor<1x2048xf32> loc(#loc5242)
    %25075 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25076 = stablehlo.reduce(%25074 init: %25075) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25077 = stablehlo.multiply %arg141, %arg141 : tensor<1x8192x2048xf32> loc(#loc5238)
    %25078 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25079 = stablehlo.reduce(%25077 init: %25078) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25080 = stablehlo.multiply %arg142, %arg142 : tensor<1x2048xf32> loc(#loc5242)
    %25081 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25082 = stablehlo.reduce(%25080 init: %25081) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25083 = stablehlo.multiply %arg143, %arg143 : tensor<1x2048xf32> loc(#loc5242)
    %25084 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25085 = stablehlo.reduce(%25083 init: %25084) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25086 = stablehlo.multiply %arg144, %arg144 : tensor<1x2048xf32> loc(#loc5242)
    %25087 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25088 = stablehlo.reduce(%25086 init: %25087) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25089 = stablehlo.multiply %arg145, %arg145 : tensor<1x2048xf32> loc(#loc5242)
    %25090 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25091 = stablehlo.reduce(%25089 init: %25090) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25092 = stablehlo.multiply %arg146, %arg146 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %25093 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25094 = stablehlo.reduce(%25092 init: %25093) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25095 = stablehlo.multiply %arg147, %arg147 : tensor<1x64xf32> loc(#loc5238)
    %25096 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25097 = stablehlo.reduce(%25095 init: %25096) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25098 = stablehlo.multiply %arg148, %arg148 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25099 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25100 = stablehlo.reduce(%25098 init: %25099) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25101 = stablehlo.multiply %arg149, %arg149 : tensor<1x8192xf32> loc(#loc5238)
    %25102 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25103 = stablehlo.reduce(%25101 init: %25102) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25104 = stablehlo.multiply %arg150, %arg150 : tensor<1x2048x8192xf32> loc(#loc5238)
    %25105 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25106 = stablehlo.reduce(%25104 init: %25105) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25107 = stablehlo.multiply %arg151, %arg151 : tensor<1x2048xf32> loc(#loc5242)
    %25108 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25109 = stablehlo.reduce(%25107 init: %25108) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25110 = stablehlo.multiply %arg152, %arg152 : tensor<1x8192x2048xf32> loc(#loc5238)
    %25111 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25112 = stablehlo.reduce(%25110 init: %25111) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25113 = stablehlo.multiply %arg153, %arg153 : tensor<1x2048xf32> loc(#loc5242)
    %25114 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25115 = stablehlo.reduce(%25113 init: %25114) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25116 = stablehlo.multiply %arg154, %arg154 : tensor<1x2048xf32> loc(#loc5242)
    %25117 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25118 = stablehlo.reduce(%25116 init: %25117) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25119 = stablehlo.multiply %arg155, %arg155 : tensor<1x2048xf32> loc(#loc5242)
    %25120 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25121 = stablehlo.reduce(%25119 init: %25120) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25122 = stablehlo.multiply %arg156, %arg156 : tensor<1x2048xf32> loc(#loc5242)
    %25123 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25124 = stablehlo.reduce(%25122 init: %25123) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25125 = stablehlo.multiply %arg157, %arg157 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %25126 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25127 = stablehlo.reduce(%25125 init: %25126) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25128 = stablehlo.multiply %arg158, %arg158 : tensor<1x64xf32> loc(#loc5238)
    %25129 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25130 = stablehlo.reduce(%25128 init: %25129) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25131 = stablehlo.multiply %arg159, %arg159 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25132 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25133 = stablehlo.reduce(%25131 init: %25132) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25134 = stablehlo.multiply %arg160, %arg160 : tensor<1x8192xf32> loc(#loc5238)
    %25135 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25136 = stablehlo.reduce(%25134 init: %25135) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25137 = stablehlo.multiply %arg161, %arg161 : tensor<1x2048x8192xf32> loc(#loc5238)
    %25138 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25139 = stablehlo.reduce(%25137 init: %25138) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25140 = stablehlo.multiply %arg162, %arg162 : tensor<1x2048xf32> loc(#loc5242)
    %25141 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25142 = stablehlo.reduce(%25140 init: %25141) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25143 = stablehlo.multiply %arg163, %arg163 : tensor<1x8192x2048xf32> loc(#loc5238)
    %25144 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25145 = stablehlo.reduce(%25143 init: %25144) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25146 = stablehlo.multiply %arg164, %arg164 : tensor<1x2048xf32> loc(#loc5242)
    %25147 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25148 = stablehlo.reduce(%25146 init: %25147) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25149 = stablehlo.multiply %arg165, %arg165 : tensor<1x2048xf32> loc(#loc5242)
    %25150 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25151 = stablehlo.reduce(%25149 init: %25150) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25152 = stablehlo.multiply %arg166, %arg166 : tensor<1x2048xf32> loc(#loc5242)
    %25153 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25154 = stablehlo.reduce(%25152 init: %25153) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25155 = stablehlo.multiply %arg167, %arg167 : tensor<1x2048xf32> loc(#loc5242)
    %25156 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25157 = stablehlo.reduce(%25155 init: %25156) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25158 = stablehlo.multiply %arg168, %arg168 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %25159 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25160 = stablehlo.reduce(%25158 init: %25159) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25161 = stablehlo.multiply %arg169, %arg169 : tensor<1x64xf32> loc(#loc5238)
    %25162 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25163 = stablehlo.reduce(%25161 init: %25162) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25164 = stablehlo.multiply %arg170, %arg170 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25165 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25166 = stablehlo.reduce(%25164 init: %25165) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25167 = stablehlo.multiply %arg171, %arg171 : tensor<1x8192xf32> loc(#loc5238)
    %25168 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25169 = stablehlo.reduce(%25167 init: %25168) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25170 = stablehlo.multiply %arg172, %arg172 : tensor<1x2048x8192xf32> loc(#loc5238)
    %25171 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25172 = stablehlo.reduce(%25170 init: %25171) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25173 = stablehlo.multiply %arg173, %arg173 : tensor<1x2048xf32> loc(#loc5242)
    %25174 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25175 = stablehlo.reduce(%25173 init: %25174) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25176 = stablehlo.multiply %arg174, %arg174 : tensor<1x8192x2048xf32> loc(#loc5238)
    %25177 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25178 = stablehlo.reduce(%25176 init: %25177) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25179 = stablehlo.multiply %arg175, %arg175 : tensor<1x2048xf32> loc(#loc5242)
    %25180 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25181 = stablehlo.reduce(%25179 init: %25180) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25182 = stablehlo.multiply %arg176, %arg176 : tensor<1x2048xf32> loc(#loc5242)
    %25183 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25184 = stablehlo.reduce(%25182 init: %25183) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25185 = stablehlo.multiply %arg177, %arg177 : tensor<1x2048xf32> loc(#loc5242)
    %25186 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25187 = stablehlo.reduce(%25185 init: %25186) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25188 = stablehlo.multiply %arg178, %arg178 : tensor<1x2048xf32> loc(#loc5242)
    %25189 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25190 = stablehlo.reduce(%25188 init: %25189) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25191 = stablehlo.multiply %arg179, %arg179 : tensor<1x3x2048x32x64xf32> loc(#loc5238)
    %25192 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25193 = stablehlo.reduce(%25191 init: %25192) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25194 = stablehlo.multiply %arg180, %arg180 : tensor<1x64xf32> loc(#loc5238)
    %25195 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25196 = stablehlo.reduce(%25194 init: %25195) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25197 = stablehlo.multiply %arg181, %arg181 : tensor<1x2048x32x64xf32> loc(#loc5238)
    %25198 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25199 = stablehlo.reduce(%25197 init: %25198) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25200 = stablehlo.broadcast_in_dim %24659, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25201 = stablehlo.broadcast_in_dim %24662, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25202 = stablehlo.broadcast_in_dim %24665, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25203 = stablehlo.broadcast_in_dim %24668, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25204 = stablehlo.broadcast_in_dim %24671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25205 = stablehlo.broadcast_in_dim %24674, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25206 = stablehlo.broadcast_in_dim %24677, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25207 = stablehlo.broadcast_in_dim %24680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25208 = stablehlo.broadcast_in_dim %24683, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25209 = stablehlo.broadcast_in_dim %24686, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25210 = stablehlo.broadcast_in_dim %24689, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25211 = stablehlo.broadcast_in_dim %24692, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25212 = stablehlo.broadcast_in_dim %24695, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25213 = stablehlo.broadcast_in_dim %24698, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25214 = stablehlo.broadcast_in_dim %24701, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25215 = stablehlo.broadcast_in_dim %24704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25216 = stablehlo.broadcast_in_dim %24707, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25217 = stablehlo.broadcast_in_dim %24710, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25218 = stablehlo.broadcast_in_dim %24713, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25219 = stablehlo.broadcast_in_dim %24716, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25220 = stablehlo.broadcast_in_dim %24719, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25221 = stablehlo.broadcast_in_dim %24722, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25222 = stablehlo.broadcast_in_dim %24725, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25223 = stablehlo.broadcast_in_dim %24728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25224 = stablehlo.broadcast_in_dim %24731, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25225 = stablehlo.broadcast_in_dim %24734, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25226 = stablehlo.broadcast_in_dim %24737, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25227 = stablehlo.broadcast_in_dim %24740, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25228 = stablehlo.broadcast_in_dim %24743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25229 = stablehlo.broadcast_in_dim %24746, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25230 = stablehlo.broadcast_in_dim %24749, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25231 = stablehlo.broadcast_in_dim %24752, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25232 = stablehlo.broadcast_in_dim %24755, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25233 = stablehlo.broadcast_in_dim %24758, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25234 = stablehlo.broadcast_in_dim %24761, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25235 = stablehlo.broadcast_in_dim %24764, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25236 = stablehlo.broadcast_in_dim %24767, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25237 = stablehlo.broadcast_in_dim %24770, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25238 = stablehlo.broadcast_in_dim %24773, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25239 = stablehlo.broadcast_in_dim %24776, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25240 = stablehlo.broadcast_in_dim %24779, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25241 = stablehlo.broadcast_in_dim %24782, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25242 = stablehlo.broadcast_in_dim %24785, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25243 = stablehlo.broadcast_in_dim %24788, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25244 = stablehlo.broadcast_in_dim %24791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25245 = stablehlo.broadcast_in_dim %24794, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25246 = stablehlo.broadcast_in_dim %24797, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25247 = stablehlo.broadcast_in_dim %24800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25248 = stablehlo.broadcast_in_dim %24803, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25249 = stablehlo.broadcast_in_dim %24806, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25250 = stablehlo.broadcast_in_dim %24809, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25251 = stablehlo.broadcast_in_dim %24812, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25252 = stablehlo.broadcast_in_dim %24815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25253 = stablehlo.broadcast_in_dim %24818, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25254 = stablehlo.broadcast_in_dim %24821, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25255 = stablehlo.broadcast_in_dim %24824, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25256 = stablehlo.broadcast_in_dim %24827, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25257 = stablehlo.broadcast_in_dim %24830, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25258 = stablehlo.broadcast_in_dim %24833, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25259 = stablehlo.broadcast_in_dim %24836, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25260 = stablehlo.broadcast_in_dim %24839, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25261 = stablehlo.broadcast_in_dim %24842, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25262 = stablehlo.broadcast_in_dim %24845, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25263 = stablehlo.broadcast_in_dim %24848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25264 = stablehlo.broadcast_in_dim %24851, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25265 = stablehlo.broadcast_in_dim %24854, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25266 = stablehlo.broadcast_in_dim %24857, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25267 = stablehlo.broadcast_in_dim %24860, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25268 = stablehlo.broadcast_in_dim %24863, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25269 = stablehlo.broadcast_in_dim %24866, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25270 = stablehlo.broadcast_in_dim %24869, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25271 = stablehlo.broadcast_in_dim %24872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25272 = stablehlo.broadcast_in_dim %24875, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25273 = stablehlo.broadcast_in_dim %24878, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25274 = stablehlo.broadcast_in_dim %24881, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25275 = stablehlo.broadcast_in_dim %24884, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25276 = stablehlo.broadcast_in_dim %24887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25277 = stablehlo.broadcast_in_dim %24890, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25278 = stablehlo.broadcast_in_dim %24893, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25279 = stablehlo.broadcast_in_dim %24896, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25280 = stablehlo.broadcast_in_dim %24899, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25281 = stablehlo.broadcast_in_dim %24902, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25282 = stablehlo.broadcast_in_dim %24905, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25283 = stablehlo.broadcast_in_dim %24908, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25284 = stablehlo.broadcast_in_dim %24911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25285 = stablehlo.broadcast_in_dim %24914, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25286 = stablehlo.broadcast_in_dim %24917, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25287 = stablehlo.broadcast_in_dim %24920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25288 = stablehlo.broadcast_in_dim %24923, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25289 = stablehlo.broadcast_in_dim %24926, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25290 = stablehlo.broadcast_in_dim %24929, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25291 = stablehlo.broadcast_in_dim %24932, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25292 = stablehlo.broadcast_in_dim %24935, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25293 = stablehlo.broadcast_in_dim %24938, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25294 = stablehlo.broadcast_in_dim %24941, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25295 = stablehlo.broadcast_in_dim %24944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25296 = stablehlo.broadcast_in_dim %24947, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25297 = stablehlo.broadcast_in_dim %24950, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25298 = stablehlo.broadcast_in_dim %24953, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25299 = stablehlo.broadcast_in_dim %24956, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25300 = stablehlo.broadcast_in_dim %24959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25301 = stablehlo.broadcast_in_dim %24962, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25302 = stablehlo.broadcast_in_dim %24965, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25303 = stablehlo.broadcast_in_dim %24968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25304 = stablehlo.broadcast_in_dim %24971, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25305 = stablehlo.broadcast_in_dim %24974, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25306 = stablehlo.broadcast_in_dim %24977, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25307 = stablehlo.broadcast_in_dim %24980, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25308 = stablehlo.broadcast_in_dim %24983, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25309 = stablehlo.broadcast_in_dim %24986, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25310 = stablehlo.broadcast_in_dim %24989, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25311 = stablehlo.broadcast_in_dim %24992, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25312 = stablehlo.broadcast_in_dim %24995, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25313 = stablehlo.broadcast_in_dim %24998, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25314 = stablehlo.broadcast_in_dim %25001, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25315 = stablehlo.broadcast_in_dim %25004, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25316 = stablehlo.broadcast_in_dim %25007, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25317 = stablehlo.broadcast_in_dim %25010, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25318 = stablehlo.broadcast_in_dim %25013, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25319 = stablehlo.broadcast_in_dim %25016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25320 = stablehlo.broadcast_in_dim %25019, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25321 = stablehlo.broadcast_in_dim %25022, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25322 = stablehlo.broadcast_in_dim %25025, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25323 = stablehlo.broadcast_in_dim %25028, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25324 = stablehlo.broadcast_in_dim %25031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25325 = stablehlo.broadcast_in_dim %25034, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25326 = stablehlo.broadcast_in_dim %25037, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25327 = stablehlo.broadcast_in_dim %25040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25328 = stablehlo.broadcast_in_dim %25043, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25329 = stablehlo.broadcast_in_dim %25046, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25330 = stablehlo.broadcast_in_dim %25049, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25331 = stablehlo.broadcast_in_dim %25052, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25332 = stablehlo.broadcast_in_dim %25055, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25333 = stablehlo.broadcast_in_dim %25058, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25334 = stablehlo.broadcast_in_dim %25061, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25335 = stablehlo.broadcast_in_dim %25064, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25336 = stablehlo.broadcast_in_dim %25067, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25337 = stablehlo.broadcast_in_dim %25070, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25338 = stablehlo.broadcast_in_dim %25073, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25339 = stablehlo.broadcast_in_dim %25076, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25340 = stablehlo.broadcast_in_dim %25079, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25341 = stablehlo.broadcast_in_dim %25082, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25342 = stablehlo.broadcast_in_dim %25085, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25343 = stablehlo.broadcast_in_dim %25088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25344 = stablehlo.broadcast_in_dim %25091, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25345 = stablehlo.broadcast_in_dim %25094, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25346 = stablehlo.broadcast_in_dim %25097, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25347 = stablehlo.broadcast_in_dim %25100, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25348 = stablehlo.broadcast_in_dim %25103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25349 = stablehlo.broadcast_in_dim %25106, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25350 = stablehlo.broadcast_in_dim %25109, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25351 = stablehlo.broadcast_in_dim %25112, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25352 = stablehlo.broadcast_in_dim %25115, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25353 = stablehlo.broadcast_in_dim %25118, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25354 = stablehlo.broadcast_in_dim %25121, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25355 = stablehlo.broadcast_in_dim %25124, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25356 = stablehlo.broadcast_in_dim %25127, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25357 = stablehlo.broadcast_in_dim %25130, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25358 = stablehlo.broadcast_in_dim %25133, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25359 = stablehlo.broadcast_in_dim %25136, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25360 = stablehlo.broadcast_in_dim %25139, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25361 = stablehlo.broadcast_in_dim %25142, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25362 = stablehlo.broadcast_in_dim %25145, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25363 = stablehlo.broadcast_in_dim %25148, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25364 = stablehlo.broadcast_in_dim %25151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25365 = stablehlo.broadcast_in_dim %25154, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25366 = stablehlo.broadcast_in_dim %25157, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25367 = stablehlo.broadcast_in_dim %25160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25368 = stablehlo.broadcast_in_dim %25163, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25369 = stablehlo.broadcast_in_dim %25166, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25370 = stablehlo.broadcast_in_dim %25169, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25371 = stablehlo.broadcast_in_dim %25172, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25372 = stablehlo.broadcast_in_dim %25175, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25373 = stablehlo.broadcast_in_dim %25178, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25374 = stablehlo.broadcast_in_dim %25181, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25375 = stablehlo.broadcast_in_dim %25184, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25376 = stablehlo.broadcast_in_dim %25187, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25377 = stablehlo.broadcast_in_dim %25190, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25378 = stablehlo.broadcast_in_dim %25193, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25379 = stablehlo.broadcast_in_dim %25196, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25380 = stablehlo.broadcast_in_dim %25199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5303)
    %25381 = stablehlo.concatenate %25200, %25201, %25202, %25203, %25204, %25205, %25206, %25207, %25208, %25209, %25210, %25211, %25212, %25213, %25214, %25215, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25382 = stablehlo.concatenate %25216, %25217, %25218, %25219, %25220, %25221, %25222, %25223, %25224, %25225, %25226, %25227, %25228, %25229, %25230, %25231, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25383 = stablehlo.concatenate %25232, %25233, %25234, %25235, %25236, %25237, %25238, %25239, %25240, %25241, %25242, %25243, %25244, %25245, %25246, %25247, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25384 = stablehlo.concatenate %25248, %25249, %25250, %25251, %25252, %25253, %25254, %25255, %25256, %25257, %25258, %25259, %25260, %25261, %25262, %25263, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25385 = stablehlo.concatenate %25264, %25265, %25266, %25267, %25268, %25269, %25270, %25271, %25272, %25273, %25274, %25275, %25276, %25277, %25278, %25279, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25386 = stablehlo.concatenate %25280, %25281, %25282, %25283, %25284, %25285, %25286, %25287, %25288, %25289, %25290, %25291, %25292, %25293, %25294, %25295, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25387 = stablehlo.concatenate %25296, %25297, %25298, %25299, %25300, %25301, %25302, %25303, %25304, %25305, %25306, %25307, %25308, %25309, %25310, %25311, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25388 = stablehlo.concatenate %25312, %25313, %25314, %25315, %25316, %25317, %25318, %25319, %25320, %25321, %25322, %25323, %25324, %25325, %25326, %25327, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25389 = stablehlo.concatenate %25328, %25329, %25330, %25331, %25332, %25333, %25334, %25335, %25336, %25337, %25338, %25339, %25340, %25341, %25342, %25343, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25390 = stablehlo.concatenate %25344, %25345, %25346, %25347, %25348, %25349, %25350, %25351, %25352, %25353, %25354, %25355, %25356, %25357, %25358, %25359, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25391 = stablehlo.concatenate %25360, %25361, %25362, %25363, %25364, %25365, %25366, %25367, %25368, %25369, %25370, %25371, %25372, %25373, %25374, %25375, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<16xf32> loc(#loc5304)
    %25392 = stablehlo.concatenate %25376, %25377, %25378, %25379, %25380, dim = 0 : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5xf32> loc(#loc5304)
    %25393 = stablehlo.concatenate %25381, %25382, %25383, %25384, %25385, %25386, %25387, %25388, %25389, %25390, %25391, %25392, dim = 0 : (tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<5xf32>) -> tensor<181xf32> loc(#loc5304)
    %25394 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5248)
    %25395 = stablehlo.reduce(%25393 init: %25394) across dimensions = [0] : (tensor<181xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5248)
      stablehlo.return %29929 : tensor<f32> loc(#loc5248)
    } loc(#loc5248)
    %25396 = stablehlo.sqrt %25395 : tensor<f32> loc(#loc5249)
    %25397 = stablehlo.add %arg1, %23006 : tensor<2048xf32> loc(#loc5263)
    %25398 = stablehlo.add %arg2, %23009 : tensor<2048xf32> loc(#loc5263)
    %25399 = stablehlo.add %arg3, %23012 : tensor<2048x2048xf32> loc(#loc5263)
    %25400 = stablehlo.add %arg4, %23015 : tensor<51200xf32> loc(#loc5263)
    %25401 = stablehlo.add %arg5, %23018 : tensor<2048x51200xf32> loc(#loc5263)
    %25402 = stablehlo.add %arg6, %23021 : tensor<1x8192xf32> loc(#loc5305)
    %25403 = stablehlo.add %arg7, %23024 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25404 = stablehlo.add %arg8, %23027 : tensor<1x2048xf32> loc(#loc5305)
    %25405 = stablehlo.add %arg9, %23030 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25406 = stablehlo.add %arg10, %23033 : tensor<1x2048xf32> loc(#loc5305)
    %25407 = stablehlo.add %arg11, %23036 : tensor<1x2048xf32> loc(#loc5305)
    %25408 = stablehlo.add %arg12, %23039 : tensor<1x2048xf32> loc(#loc5305)
    %25409 = stablehlo.add %arg13, %23042 : tensor<1x2048xf32> loc(#loc5305)
    %25410 = stablehlo.add %arg14, %23045 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25411 = stablehlo.add %arg15, %23048 : tensor<1x64xf32> loc(#loc5305)
    %25412 = stablehlo.add %arg16, %23051 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25413 = stablehlo.add %arg17, %23054 : tensor<1x8192xf32> loc(#loc5305)
    %25414 = stablehlo.add %arg18, %23057 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25415 = stablehlo.add %arg19, %23060 : tensor<1x2048xf32> loc(#loc5305)
    %25416 = stablehlo.add %arg20, %23063 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25417 = stablehlo.add %arg21, %23066 : tensor<1x2048xf32> loc(#loc5305)
    %25418 = stablehlo.add %arg22, %23069 : tensor<1x2048xf32> loc(#loc5305)
    %25419 = stablehlo.add %arg23, %23072 : tensor<1x2048xf32> loc(#loc5305)
    %25420 = stablehlo.add %arg24, %23075 : tensor<1x2048xf32> loc(#loc5305)
    %25421 = stablehlo.add %arg25, %23078 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25422 = stablehlo.add %arg26, %23081 : tensor<1x64xf32> loc(#loc5305)
    %25423 = stablehlo.add %arg27, %23084 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25424 = stablehlo.add %arg28, %23087 : tensor<1x8192xf32> loc(#loc5305)
    %25425 = stablehlo.add %arg29, %23090 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25426 = stablehlo.add %arg30, %23093 : tensor<1x2048xf32> loc(#loc5305)
    %25427 = stablehlo.add %arg31, %23096 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25428 = stablehlo.add %arg32, %23099 : tensor<1x2048xf32> loc(#loc5305)
    %25429 = stablehlo.add %arg33, %23102 : tensor<1x2048xf32> loc(#loc5305)
    %25430 = stablehlo.add %arg34, %23105 : tensor<1x2048xf32> loc(#loc5305)
    %25431 = stablehlo.add %arg35, %23108 : tensor<1x2048xf32> loc(#loc5305)
    %25432 = stablehlo.add %arg36, %23111 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25433 = stablehlo.add %arg37, %23114 : tensor<1x64xf32> loc(#loc5305)
    %25434 = stablehlo.add %arg38, %23117 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25435 = stablehlo.add %arg39, %23120 : tensor<1x8192xf32> loc(#loc5305)
    %25436 = stablehlo.add %arg40, %23123 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25437 = stablehlo.add %arg41, %23126 : tensor<1x2048xf32> loc(#loc5305)
    %25438 = stablehlo.add %arg42, %23129 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25439 = stablehlo.add %arg43, %23132 : tensor<1x2048xf32> loc(#loc5305)
    %25440 = stablehlo.add %arg44, %23135 : tensor<1x2048xf32> loc(#loc5305)
    %25441 = stablehlo.add %arg45, %23138 : tensor<1x2048xf32> loc(#loc5305)
    %25442 = stablehlo.add %arg46, %23141 : tensor<1x2048xf32> loc(#loc5305)
    %25443 = stablehlo.add %arg47, %23144 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25444 = stablehlo.add %arg48, %23147 : tensor<1x64xf32> loc(#loc5305)
    %25445 = stablehlo.add %arg49, %23150 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25446 = stablehlo.add %arg50, %23153 : tensor<1x8192xf32> loc(#loc5305)
    %25447 = stablehlo.add %arg51, %23156 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25448 = stablehlo.add %arg52, %23159 : tensor<1x2048xf32> loc(#loc5305)
    %25449 = stablehlo.add %arg53, %23162 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25450 = stablehlo.add %arg54, %23165 : tensor<1x2048xf32> loc(#loc5305)
    %25451 = stablehlo.add %arg55, %23168 : tensor<1x2048xf32> loc(#loc5305)
    %25452 = stablehlo.add %arg56, %23171 : tensor<1x2048xf32> loc(#loc5305)
    %25453 = stablehlo.add %arg57, %23174 : tensor<1x2048xf32> loc(#loc5305)
    %25454 = stablehlo.add %arg58, %23177 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25455 = stablehlo.add %arg59, %23180 : tensor<1x64xf32> loc(#loc5305)
    %25456 = stablehlo.add %arg60, %23183 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25457 = stablehlo.add %arg61, %23186 : tensor<1x8192xf32> loc(#loc5305)
    %25458 = stablehlo.add %arg62, %23189 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25459 = stablehlo.add %arg63, %23192 : tensor<1x2048xf32> loc(#loc5305)
    %25460 = stablehlo.add %arg64, %23195 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25461 = stablehlo.add %arg65, %23198 : tensor<1x2048xf32> loc(#loc5305)
    %25462 = stablehlo.add %arg66, %23201 : tensor<1x2048xf32> loc(#loc5305)
    %25463 = stablehlo.add %arg67, %23204 : tensor<1x2048xf32> loc(#loc5305)
    %25464 = stablehlo.add %arg68, %23207 : tensor<1x2048xf32> loc(#loc5305)
    %25465 = stablehlo.add %arg69, %23210 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25466 = stablehlo.add %arg70, %23213 : tensor<1x64xf32> loc(#loc5305)
    %25467 = stablehlo.add %arg71, %23216 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25468 = stablehlo.add %arg72, %23219 : tensor<1x8192xf32> loc(#loc5305)
    %25469 = stablehlo.add %arg73, %23222 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25470 = stablehlo.add %arg74, %23225 : tensor<1x2048xf32> loc(#loc5305)
    %25471 = stablehlo.add %arg75, %23228 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25472 = stablehlo.add %arg76, %23231 : tensor<1x2048xf32> loc(#loc5305)
    %25473 = stablehlo.add %arg77, %23234 : tensor<1x2048xf32> loc(#loc5305)
    %25474 = stablehlo.add %arg78, %23237 : tensor<1x2048xf32> loc(#loc5305)
    %25475 = stablehlo.add %arg79, %23240 : tensor<1x2048xf32> loc(#loc5305)
    %25476 = stablehlo.add %arg80, %23243 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25477 = stablehlo.add %arg81, %23246 : tensor<1x64xf32> loc(#loc5305)
    %25478 = stablehlo.add %arg82, %23249 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25479 = stablehlo.add %arg83, %23252 : tensor<1x8192xf32> loc(#loc5305)
    %25480 = stablehlo.add %arg84, %23255 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25481 = stablehlo.add %arg85, %23258 : tensor<1x2048xf32> loc(#loc5305)
    %25482 = stablehlo.add %arg86, %23261 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25483 = stablehlo.add %arg87, %23264 : tensor<1x2048xf32> loc(#loc5305)
    %25484 = stablehlo.add %arg88, %23267 : tensor<1x2048xf32> loc(#loc5305)
    %25485 = stablehlo.add %arg89, %23270 : tensor<1x2048xf32> loc(#loc5305)
    %25486 = stablehlo.add %arg90, %23273 : tensor<1x2048xf32> loc(#loc5305)
    %25487 = stablehlo.add %arg91, %23276 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25488 = stablehlo.add %arg92, %23279 : tensor<1x64xf32> loc(#loc5305)
    %25489 = stablehlo.add %arg93, %23282 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25490 = stablehlo.add %arg94, %23285 : tensor<1x8192xf32> loc(#loc5305)
    %25491 = stablehlo.add %arg95, %23288 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25492 = stablehlo.add %arg96, %23291 : tensor<1x2048xf32> loc(#loc5305)
    %25493 = stablehlo.add %arg97, %23294 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25494 = stablehlo.add %arg98, %23297 : tensor<1x2048xf32> loc(#loc5305)
    %25495 = stablehlo.add %arg99, %23300 : tensor<1x2048xf32> loc(#loc5305)
    %25496 = stablehlo.add %arg100, %23303 : tensor<1x2048xf32> loc(#loc5305)
    %25497 = stablehlo.add %arg101, %23306 : tensor<1x2048xf32> loc(#loc5305)
    %25498 = stablehlo.add %arg102, %23309 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25499 = stablehlo.add %arg103, %23312 : tensor<1x64xf32> loc(#loc5305)
    %25500 = stablehlo.add %arg104, %23315 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25501 = stablehlo.add %arg105, %23318 : tensor<1x8192xf32> loc(#loc5305)
    %25502 = stablehlo.add %arg106, %23321 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25503 = stablehlo.add %arg107, %23324 : tensor<1x2048xf32> loc(#loc5305)
    %25504 = stablehlo.add %arg108, %23327 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25505 = stablehlo.add %arg109, %23330 : tensor<1x2048xf32> loc(#loc5305)
    %25506 = stablehlo.add %arg110, %23333 : tensor<1x2048xf32> loc(#loc5305)
    %25507 = stablehlo.add %arg111, %23336 : tensor<1x2048xf32> loc(#loc5305)
    %25508 = stablehlo.add %arg112, %23339 : tensor<1x2048xf32> loc(#loc5305)
    %25509 = stablehlo.add %arg113, %23342 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25510 = stablehlo.add %arg114, %23345 : tensor<1x64xf32> loc(#loc5305)
    %25511 = stablehlo.add %arg115, %23348 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25512 = stablehlo.add %arg116, %23351 : tensor<1x8192xf32> loc(#loc5305)
    %25513 = stablehlo.add %arg117, %23354 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25514 = stablehlo.add %arg118, %23357 : tensor<1x2048xf32> loc(#loc5305)
    %25515 = stablehlo.add %arg119, %23360 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25516 = stablehlo.add %arg120, %23363 : tensor<1x2048xf32> loc(#loc5305)
    %25517 = stablehlo.add %arg121, %23366 : tensor<1x2048xf32> loc(#loc5305)
    %25518 = stablehlo.add %arg122, %23369 : tensor<1x2048xf32> loc(#loc5305)
    %25519 = stablehlo.add %arg123, %23372 : tensor<1x2048xf32> loc(#loc5305)
    %25520 = stablehlo.add %arg124, %23375 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25521 = stablehlo.add %arg125, %23378 : tensor<1x64xf32> loc(#loc5305)
    %25522 = stablehlo.add %arg126, %23381 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25523 = stablehlo.add %arg127, %23384 : tensor<1x8192xf32> loc(#loc5305)
    %25524 = stablehlo.add %arg128, %23387 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25525 = stablehlo.add %arg129, %23390 : tensor<1x2048xf32> loc(#loc5305)
    %25526 = stablehlo.add %arg130, %23393 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25527 = stablehlo.add %arg131, %23396 : tensor<1x2048xf32> loc(#loc5305)
    %25528 = stablehlo.add %arg132, %23399 : tensor<1x2048xf32> loc(#loc5305)
    %25529 = stablehlo.add %arg133, %23402 : tensor<1x2048xf32> loc(#loc5305)
    %25530 = stablehlo.add %arg134, %23405 : tensor<1x2048xf32> loc(#loc5305)
    %25531 = stablehlo.add %arg135, %23408 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25532 = stablehlo.add %arg136, %23411 : tensor<1x64xf32> loc(#loc5305)
    %25533 = stablehlo.add %arg137, %23414 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25534 = stablehlo.add %arg138, %23417 : tensor<1x8192xf32> loc(#loc5305)
    %25535 = stablehlo.add %arg139, %23420 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25536 = stablehlo.add %arg140, %23423 : tensor<1x2048xf32> loc(#loc5305)
    %25537 = stablehlo.add %arg141, %23426 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25538 = stablehlo.add %arg142, %23429 : tensor<1x2048xf32> loc(#loc5305)
    %25539 = stablehlo.add %arg143, %23432 : tensor<1x2048xf32> loc(#loc5305)
    %25540 = stablehlo.add %arg144, %23435 : tensor<1x2048xf32> loc(#loc5305)
    %25541 = stablehlo.add %arg145, %23438 : tensor<1x2048xf32> loc(#loc5305)
    %25542 = stablehlo.add %arg146, %23441 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25543 = stablehlo.add %arg147, %23444 : tensor<1x64xf32> loc(#loc5305)
    %25544 = stablehlo.add %arg148, %23447 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25545 = stablehlo.add %arg149, %23450 : tensor<1x8192xf32> loc(#loc5305)
    %25546 = stablehlo.add %arg150, %23453 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25547 = stablehlo.add %arg151, %23456 : tensor<1x2048xf32> loc(#loc5305)
    %25548 = stablehlo.add %arg152, %23459 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25549 = stablehlo.add %arg153, %23462 : tensor<1x2048xf32> loc(#loc5305)
    %25550 = stablehlo.add %arg154, %23465 : tensor<1x2048xf32> loc(#loc5305)
    %25551 = stablehlo.add %arg155, %23468 : tensor<1x2048xf32> loc(#loc5305)
    %25552 = stablehlo.add %arg156, %23471 : tensor<1x2048xf32> loc(#loc5305)
    %25553 = stablehlo.add %arg157, %23474 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25554 = stablehlo.add %arg158, %23477 : tensor<1x64xf32> loc(#loc5305)
    %25555 = stablehlo.add %arg159, %23480 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25556 = stablehlo.add %arg160, %23483 : tensor<1x8192xf32> loc(#loc5305)
    %25557 = stablehlo.add %arg161, %23486 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25558 = stablehlo.add %arg162, %23489 : tensor<1x2048xf32> loc(#loc5305)
    %25559 = stablehlo.add %arg163, %23492 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25560 = stablehlo.add %arg164, %23495 : tensor<1x2048xf32> loc(#loc5305)
    %25561 = stablehlo.add %arg165, %23498 : tensor<1x2048xf32> loc(#loc5305)
    %25562 = stablehlo.add %arg166, %23501 : tensor<1x2048xf32> loc(#loc5305)
    %25563 = stablehlo.add %arg167, %23504 : tensor<1x2048xf32> loc(#loc5305)
    %25564 = stablehlo.add %arg168, %23507 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25565 = stablehlo.add %arg169, %23510 : tensor<1x64xf32> loc(#loc5305)
    %25566 = stablehlo.add %arg170, %23513 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25567 = stablehlo.add %arg171, %23516 : tensor<1x8192xf32> loc(#loc5305)
    %25568 = stablehlo.add %arg172, %23519 : tensor<1x2048x8192xf32> loc(#loc5305)
    %25569 = stablehlo.add %arg173, %23522 : tensor<1x2048xf32> loc(#loc5305)
    %25570 = stablehlo.add %arg174, %23525 : tensor<1x8192x2048xf32> loc(#loc5305)
    %25571 = stablehlo.add %arg175, %23528 : tensor<1x2048xf32> loc(#loc5305)
    %25572 = stablehlo.add %arg176, %23531 : tensor<1x2048xf32> loc(#loc5305)
    %25573 = stablehlo.add %arg177, %23534 : tensor<1x2048xf32> loc(#loc5305)
    %25574 = stablehlo.add %arg178, %23537 : tensor<1x2048xf32> loc(#loc5305)
    %25575 = stablehlo.add %arg179, %23540 : tensor<1x3x2048x32x64xf32> loc(#loc5305)
    %25576 = stablehlo.add %arg180, %23543 : tensor<1x64xf32> loc(#loc5305)
    %25577 = stablehlo.add %arg181, %23546 : tensor<1x2048x32x64xf32> loc(#loc5305)
    %25578 = stablehlo.constant dense<1> : tensor<ui32> loc(#loc)
    %25579 = stablehlo.add %arg0, %25578 : tensor<ui32> loc(#loc5306)
    %25580 = stablehlo.abs %25397 : tensor<2048xf32> loc(#loc5307)
    %25581 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5308)
    %25582 = stablehlo.reduce(%25580 init: %25581) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5308)
      stablehlo.return %29929 : tensor<f32> loc(#loc5308)
    } loc(#loc5308)
    %25583 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25584 = stablehlo.maximum %25582, %25583 : tensor<f32> loc(#loc5309)
    %25585 = stablehlo.broadcast_in_dim %25584, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5272)
    %25586 = stablehlo.divide %25397, %25585 : tensor<2048xf32> loc(#loc5272)
    %25587 = stablehlo.multiply %25584, %25584 : tensor<f32> loc(#loc5310)
    %25588 = stablehlo.multiply %25586, %25586 : tensor<2048xf32> loc(#loc5310)
    %25589 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %25590 = stablehlo.reduce(%25588 init: %25589) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %25591 = stablehlo.multiply %25587, %25590 : tensor<f32> loc(#loc5310)
    %25592 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25593 = stablehlo.broadcast_in_dim %25592, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25594 = stablehlo.broadcast_in_dim %25591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25595 = stablehlo.concatenate %25593, %25594, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25596 = stablehlo.abs %25398 : tensor<2048xf32> loc(#loc5307)
    %25597 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5308)
    %25598 = stablehlo.reduce(%25596 init: %25597) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5308)
      stablehlo.return %29929 : tensor<f32> loc(#loc5308)
    } loc(#loc5308)
    %25599 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25600 = stablehlo.maximum %25598, %25599 : tensor<f32> loc(#loc5309)
    %25601 = stablehlo.broadcast_in_dim %25600, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc5272)
    %25602 = stablehlo.divide %25398, %25601 : tensor<2048xf32> loc(#loc5272)
    %25603 = stablehlo.multiply %25600, %25600 : tensor<f32> loc(#loc5310)
    %25604 = stablehlo.multiply %25602, %25602 : tensor<2048xf32> loc(#loc5310)
    %25605 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %25606 = stablehlo.reduce(%25604 init: %25605) across dimensions = [0] : (tensor<2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %25607 = stablehlo.multiply %25603, %25606 : tensor<f32> loc(#loc5310)
    %25608 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25609 = stablehlo.broadcast_in_dim %25608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25610 = stablehlo.broadcast_in_dim %25607, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25611 = stablehlo.concatenate %25609, %25610, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25612 = stablehlo.abs %25399 : tensor<2048x2048xf32> loc(#loc5307)
    %25613 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25614 = stablehlo.reduce(%25612 init: %25613) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25615 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25616 = stablehlo.maximum %25614, %25615 : tensor<f32> loc(#loc5309)
    %25617 = stablehlo.broadcast_in_dim %25616, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc5272)
    %25618 = stablehlo.divide %25399, %25617 : tensor<2048x2048xf32> loc(#loc5272)
    %25619 = stablehlo.multiply %25616, %25616 : tensor<f32> loc(#loc5310)
    %25620 = stablehlo.multiply %25618, %25618 : tensor<2048x2048xf32> loc(#loc5310)
    %25621 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25622 = stablehlo.reduce(%25620 init: %25621) across dimensions = [0, 1] : (tensor<2048x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25623 = stablehlo.multiply %25619, %25622 : tensor<f32> loc(#loc5310)
    %25624 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %25625 = stablehlo.broadcast_in_dim %25624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25626 = stablehlo.broadcast_in_dim %25623, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25627 = stablehlo.concatenate %25625, %25626, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25628 = stablehlo.abs %25400 : tensor<51200xf32> loc(#loc5307)
    %25629 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5308)
    %25630 = stablehlo.reduce(%25628 init: %25629) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5308)
      stablehlo.return %29929 : tensor<f32> loc(#loc5308)
    } loc(#loc5308)
    %25631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25632 = stablehlo.maximum %25630, %25631 : tensor<f32> loc(#loc5309)
    %25633 = stablehlo.broadcast_in_dim %25632, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc5272)
    %25634 = stablehlo.divide %25400, %25633 : tensor<51200xf32> loc(#loc5272)
    %25635 = stablehlo.multiply %25632, %25632 : tensor<f32> loc(#loc5310)
    %25636 = stablehlo.multiply %25634, %25634 : tensor<51200xf32> loc(#loc5310)
    %25637 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5239)
    %25638 = stablehlo.reduce(%25636 init: %25637) across dimensions = [0] : (tensor<51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5239)
      stablehlo.return %29929 : tensor<f32> loc(#loc5239)
    } loc(#loc5239)
    %25639 = stablehlo.multiply %25635, %25638 : tensor<f32> loc(#loc5310)
    %25640 = stablehlo.constant dense<5.120000e+04> : tensor<f32> loc(#loc)
    %25641 = stablehlo.broadcast_in_dim %25640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25642 = stablehlo.broadcast_in_dim %25639, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25643 = stablehlo.concatenate %25641, %25642, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25644 = stablehlo.abs %25401 : tensor<2048x51200xf32> loc(#loc5307)
    %25645 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25646 = stablehlo.reduce(%25644 init: %25645) across dimensions = [0, 1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25647 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25648 = stablehlo.maximum %25646, %25647 : tensor<f32> loc(#loc5309)
    %25649 = stablehlo.broadcast_in_dim %25648, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc5272)
    %25650 = stablehlo.divide %25401, %25649 : tensor<2048x51200xf32> loc(#loc5272)
    %25651 = stablehlo.multiply %25648, %25648 : tensor<f32> loc(#loc5310)
    %25652 = stablehlo.multiply %25650, %25650 : tensor<2048x51200xf32> loc(#loc5310)
    %25653 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25654 = stablehlo.reduce(%25652 init: %25653) across dimensions = [0, 1] : (tensor<2048x51200xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25655 = stablehlo.multiply %25651, %25654 : tensor<f32> loc(#loc5310)
    %25656 = stablehlo.constant dense<0x4CC80000> : tensor<f32> loc(#loc)
    %25657 = stablehlo.broadcast_in_dim %25656, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25658 = stablehlo.broadcast_in_dim %25655, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25659 = stablehlo.concatenate %25657, %25658, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25660 = stablehlo.abs %25402 : tensor<1x8192xf32> loc(#loc5307)
    %25661 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25662 = stablehlo.reduce(%25660 init: %25661) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25663 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25664 = stablehlo.maximum %25662, %25663 : tensor<f32> loc(#loc5309)
    %25665 = stablehlo.broadcast_in_dim %25664, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %25666 = stablehlo.divide %25402, %25665 : tensor<1x8192xf32> loc(#loc5314)
    %25667 = stablehlo.multiply %25664, %25664 : tensor<f32> loc(#loc5310)
    %25668 = stablehlo.multiply %25666, %25666 : tensor<1x8192xf32> loc(#loc5310)
    %25669 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25670 = stablehlo.reduce(%25668 init: %25669) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25671 = stablehlo.multiply %25667, %25670 : tensor<f32> loc(#loc5310)
    %25672 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %25673 = stablehlo.broadcast_in_dim %25672, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25674 = stablehlo.broadcast_in_dim %25671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25675 = stablehlo.concatenate %25673, %25674, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25676 = stablehlo.abs %25403 : tensor<1x2048x8192xf32> loc(#loc5307)
    %25677 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %25678 = stablehlo.reduce(%25676 init: %25677) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %25679 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25680 = stablehlo.maximum %25678, %25679 : tensor<f32> loc(#loc5309)
    %25681 = stablehlo.broadcast_in_dim %25680, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %25682 = stablehlo.divide %25403, %25681 : tensor<1x2048x8192xf32> loc(#loc5314)
    %25683 = stablehlo.multiply %25680, %25680 : tensor<f32> loc(#loc5310)
    %25684 = stablehlo.multiply %25682, %25682 : tensor<1x2048x8192xf32> loc(#loc5310)
    %25685 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25686 = stablehlo.reduce(%25684 init: %25685) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25687 = stablehlo.multiply %25683, %25686 : tensor<f32> loc(#loc5310)
    %25688 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %25689 = stablehlo.broadcast_in_dim %25688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25690 = stablehlo.broadcast_in_dim %25687, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25691 = stablehlo.concatenate %25689, %25690, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25692 = stablehlo.abs %25404 : tensor<1x2048xf32> loc(#loc5307)
    %25693 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25694 = stablehlo.reduce(%25692 init: %25693) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25695 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25696 = stablehlo.maximum %25694, %25695 : tensor<f32> loc(#loc5309)
    %25697 = stablehlo.broadcast_in_dim %25696, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25698 = stablehlo.divide %25404, %25697 : tensor<1x2048xf32> loc(#loc5314)
    %25699 = stablehlo.multiply %25696, %25696 : tensor<f32> loc(#loc5310)
    %25700 = stablehlo.multiply %25698, %25698 : tensor<1x2048xf32> loc(#loc5310)
    %25701 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25702 = stablehlo.reduce(%25700 init: %25701) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25703 = stablehlo.multiply %25699, %25702 : tensor<f32> loc(#loc5310)
    %25704 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25705 = stablehlo.broadcast_in_dim %25704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25706 = stablehlo.broadcast_in_dim %25703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25707 = stablehlo.concatenate %25705, %25706, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25708 = stablehlo.abs %25405 : tensor<1x8192x2048xf32> loc(#loc5307)
    %25709 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %25710 = stablehlo.reduce(%25708 init: %25709) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %25711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25712 = stablehlo.maximum %25710, %25711 : tensor<f32> loc(#loc5309)
    %25713 = stablehlo.broadcast_in_dim %25712, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %25714 = stablehlo.divide %25405, %25713 : tensor<1x8192x2048xf32> loc(#loc5314)
    %25715 = stablehlo.multiply %25712, %25712 : tensor<f32> loc(#loc5310)
    %25716 = stablehlo.multiply %25714, %25714 : tensor<1x8192x2048xf32> loc(#loc5310)
    %25717 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25718 = stablehlo.reduce(%25716 init: %25717) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25719 = stablehlo.multiply %25715, %25718 : tensor<f32> loc(#loc5310)
    %25720 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %25721 = stablehlo.broadcast_in_dim %25720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25722 = stablehlo.broadcast_in_dim %25719, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25723 = stablehlo.concatenate %25721, %25722, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25724 = stablehlo.abs %25406 : tensor<1x2048xf32> loc(#loc5307)
    %25725 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25726 = stablehlo.reduce(%25724 init: %25725) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25727 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25728 = stablehlo.maximum %25726, %25727 : tensor<f32> loc(#loc5309)
    %25729 = stablehlo.broadcast_in_dim %25728, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25730 = stablehlo.divide %25406, %25729 : tensor<1x2048xf32> loc(#loc5314)
    %25731 = stablehlo.multiply %25728, %25728 : tensor<f32> loc(#loc5310)
    %25732 = stablehlo.multiply %25730, %25730 : tensor<1x2048xf32> loc(#loc5310)
    %25733 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25734 = stablehlo.reduce(%25732 init: %25733) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25735 = stablehlo.multiply %25731, %25734 : tensor<f32> loc(#loc5310)
    %25736 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25737 = stablehlo.broadcast_in_dim %25736, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25738 = stablehlo.broadcast_in_dim %25735, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25739 = stablehlo.concatenate %25737, %25738, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25740 = stablehlo.abs %25407 : tensor<1x2048xf32> loc(#loc5307)
    %25741 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25742 = stablehlo.reduce(%25740 init: %25741) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25744 = stablehlo.maximum %25742, %25743 : tensor<f32> loc(#loc5309)
    %25745 = stablehlo.broadcast_in_dim %25744, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25746 = stablehlo.divide %25407, %25745 : tensor<1x2048xf32> loc(#loc5314)
    %25747 = stablehlo.multiply %25744, %25744 : tensor<f32> loc(#loc5310)
    %25748 = stablehlo.multiply %25746, %25746 : tensor<1x2048xf32> loc(#loc5310)
    %25749 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25750 = stablehlo.reduce(%25748 init: %25749) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25751 = stablehlo.multiply %25747, %25750 : tensor<f32> loc(#loc5310)
    %25752 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25753 = stablehlo.broadcast_in_dim %25752, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25754 = stablehlo.broadcast_in_dim %25751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25755 = stablehlo.concatenate %25753, %25754, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25756 = stablehlo.abs %25408 : tensor<1x2048xf32> loc(#loc5307)
    %25757 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25758 = stablehlo.reduce(%25756 init: %25757) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25759 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25760 = stablehlo.maximum %25758, %25759 : tensor<f32> loc(#loc5309)
    %25761 = stablehlo.broadcast_in_dim %25760, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25762 = stablehlo.divide %25408, %25761 : tensor<1x2048xf32> loc(#loc5314)
    %25763 = stablehlo.multiply %25760, %25760 : tensor<f32> loc(#loc5310)
    %25764 = stablehlo.multiply %25762, %25762 : tensor<1x2048xf32> loc(#loc5310)
    %25765 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25766 = stablehlo.reduce(%25764 init: %25765) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25767 = stablehlo.multiply %25763, %25766 : tensor<f32> loc(#loc5310)
    %25768 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25769 = stablehlo.broadcast_in_dim %25768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25770 = stablehlo.broadcast_in_dim %25767, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25771 = stablehlo.concatenate %25769, %25770, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25772 = stablehlo.abs %25409 : tensor<1x2048xf32> loc(#loc5307)
    %25773 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25774 = stablehlo.reduce(%25772 init: %25773) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25775 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25776 = stablehlo.maximum %25774, %25775 : tensor<f32> loc(#loc5309)
    %25777 = stablehlo.broadcast_in_dim %25776, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25778 = stablehlo.divide %25409, %25777 : tensor<1x2048xf32> loc(#loc5314)
    %25779 = stablehlo.multiply %25776, %25776 : tensor<f32> loc(#loc5310)
    %25780 = stablehlo.multiply %25778, %25778 : tensor<1x2048xf32> loc(#loc5310)
    %25781 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25782 = stablehlo.reduce(%25780 init: %25781) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25783 = stablehlo.multiply %25779, %25782 : tensor<f32> loc(#loc5310)
    %25784 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25785 = stablehlo.broadcast_in_dim %25784, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25786 = stablehlo.broadcast_in_dim %25783, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25787 = stablehlo.concatenate %25785, %25786, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25788 = stablehlo.abs %25410 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %25789 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %25790 = stablehlo.reduce(%25788 init: %25789) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %25791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25792 = stablehlo.maximum %25790, %25791 : tensor<f32> loc(#loc5309)
    %25793 = stablehlo.broadcast_in_dim %25792, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %25794 = stablehlo.divide %25410, %25793 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %25795 = stablehlo.multiply %25792, %25792 : tensor<f32> loc(#loc5310)
    %25796 = stablehlo.multiply %25794, %25794 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %25797 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25798 = stablehlo.reduce(%25796 init: %25797) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25799 = stablehlo.multiply %25795, %25798 : tensor<f32> loc(#loc5310)
    %25800 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %25801 = stablehlo.broadcast_in_dim %25800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25802 = stablehlo.broadcast_in_dim %25799, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25803 = stablehlo.concatenate %25801, %25802, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25804 = stablehlo.abs %25411 : tensor<1x64xf32> loc(#loc5307)
    %25805 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25806 = stablehlo.reduce(%25804 init: %25805) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25807 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25808 = stablehlo.maximum %25806, %25807 : tensor<f32> loc(#loc5309)
    %25809 = stablehlo.broadcast_in_dim %25808, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %25810 = stablehlo.divide %25411, %25809 : tensor<1x64xf32> loc(#loc5314)
    %25811 = stablehlo.multiply %25808, %25808 : tensor<f32> loc(#loc5310)
    %25812 = stablehlo.multiply %25810, %25810 : tensor<1x64xf32> loc(#loc5310)
    %25813 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25814 = stablehlo.reduce(%25812 init: %25813) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25815 = stablehlo.multiply %25811, %25814 : tensor<f32> loc(#loc5310)
    %25816 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %25817 = stablehlo.broadcast_in_dim %25816, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25818 = stablehlo.broadcast_in_dim %25815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25819 = stablehlo.concatenate %25817, %25818, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25820 = stablehlo.abs %25412 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %25821 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %25822 = stablehlo.reduce(%25820 init: %25821) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %25823 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25824 = stablehlo.maximum %25822, %25823 : tensor<f32> loc(#loc5309)
    %25825 = stablehlo.broadcast_in_dim %25824, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %25826 = stablehlo.divide %25412, %25825 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %25827 = stablehlo.multiply %25824, %25824 : tensor<f32> loc(#loc5310)
    %25828 = stablehlo.multiply %25826, %25826 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %25829 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %25830 = stablehlo.reduce(%25828 init: %25829) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %25831 = stablehlo.multiply %25827, %25830 : tensor<f32> loc(#loc5310)
    %25832 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %25833 = stablehlo.broadcast_in_dim %25832, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25834 = stablehlo.broadcast_in_dim %25831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25835 = stablehlo.concatenate %25833, %25834, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25836 = stablehlo.abs %25413 : tensor<1x8192xf32> loc(#loc5307)
    %25837 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25838 = stablehlo.reduce(%25836 init: %25837) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25839 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25840 = stablehlo.maximum %25838, %25839 : tensor<f32> loc(#loc5309)
    %25841 = stablehlo.broadcast_in_dim %25840, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %25842 = stablehlo.divide %25413, %25841 : tensor<1x8192xf32> loc(#loc5314)
    %25843 = stablehlo.multiply %25840, %25840 : tensor<f32> loc(#loc5310)
    %25844 = stablehlo.multiply %25842, %25842 : tensor<1x8192xf32> loc(#loc5310)
    %25845 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25846 = stablehlo.reduce(%25844 init: %25845) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25847 = stablehlo.multiply %25843, %25846 : tensor<f32> loc(#loc5310)
    %25848 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %25849 = stablehlo.broadcast_in_dim %25848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25850 = stablehlo.broadcast_in_dim %25847, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25851 = stablehlo.concatenate %25849, %25850, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25852 = stablehlo.abs %25414 : tensor<1x2048x8192xf32> loc(#loc5307)
    %25853 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %25854 = stablehlo.reduce(%25852 init: %25853) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %25855 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25856 = stablehlo.maximum %25854, %25855 : tensor<f32> loc(#loc5309)
    %25857 = stablehlo.broadcast_in_dim %25856, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %25858 = stablehlo.divide %25414, %25857 : tensor<1x2048x8192xf32> loc(#loc5314)
    %25859 = stablehlo.multiply %25856, %25856 : tensor<f32> loc(#loc5310)
    %25860 = stablehlo.multiply %25858, %25858 : tensor<1x2048x8192xf32> loc(#loc5310)
    %25861 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25862 = stablehlo.reduce(%25860 init: %25861) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25863 = stablehlo.multiply %25859, %25862 : tensor<f32> loc(#loc5310)
    %25864 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %25865 = stablehlo.broadcast_in_dim %25864, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25866 = stablehlo.broadcast_in_dim %25863, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25867 = stablehlo.concatenate %25865, %25866, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25868 = stablehlo.abs %25415 : tensor<1x2048xf32> loc(#loc5307)
    %25869 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25870 = stablehlo.reduce(%25868 init: %25869) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25871 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25872 = stablehlo.maximum %25870, %25871 : tensor<f32> loc(#loc5309)
    %25873 = stablehlo.broadcast_in_dim %25872, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25874 = stablehlo.divide %25415, %25873 : tensor<1x2048xf32> loc(#loc5314)
    %25875 = stablehlo.multiply %25872, %25872 : tensor<f32> loc(#loc5310)
    %25876 = stablehlo.multiply %25874, %25874 : tensor<1x2048xf32> loc(#loc5310)
    %25877 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25878 = stablehlo.reduce(%25876 init: %25877) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25879 = stablehlo.multiply %25875, %25878 : tensor<f32> loc(#loc5310)
    %25880 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25881 = stablehlo.broadcast_in_dim %25880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25882 = stablehlo.broadcast_in_dim %25879, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25883 = stablehlo.concatenate %25881, %25882, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25884 = stablehlo.abs %25416 : tensor<1x8192x2048xf32> loc(#loc5307)
    %25885 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %25886 = stablehlo.reduce(%25884 init: %25885) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %25887 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25888 = stablehlo.maximum %25886, %25887 : tensor<f32> loc(#loc5309)
    %25889 = stablehlo.broadcast_in_dim %25888, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %25890 = stablehlo.divide %25416, %25889 : tensor<1x8192x2048xf32> loc(#loc5314)
    %25891 = stablehlo.multiply %25888, %25888 : tensor<f32> loc(#loc5310)
    %25892 = stablehlo.multiply %25890, %25890 : tensor<1x8192x2048xf32> loc(#loc5310)
    %25893 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %25894 = stablehlo.reduce(%25892 init: %25893) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %25895 = stablehlo.multiply %25891, %25894 : tensor<f32> loc(#loc5310)
    %25896 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %25897 = stablehlo.broadcast_in_dim %25896, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25898 = stablehlo.broadcast_in_dim %25895, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25899 = stablehlo.concatenate %25897, %25898, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25900 = stablehlo.abs %25417 : tensor<1x2048xf32> loc(#loc5307)
    %25901 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25902 = stablehlo.reduce(%25900 init: %25901) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25903 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25904 = stablehlo.maximum %25902, %25903 : tensor<f32> loc(#loc5309)
    %25905 = stablehlo.broadcast_in_dim %25904, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25906 = stablehlo.divide %25417, %25905 : tensor<1x2048xf32> loc(#loc5314)
    %25907 = stablehlo.multiply %25904, %25904 : tensor<f32> loc(#loc5310)
    %25908 = stablehlo.multiply %25906, %25906 : tensor<1x2048xf32> loc(#loc5310)
    %25909 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25910 = stablehlo.reduce(%25908 init: %25909) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25911 = stablehlo.multiply %25907, %25910 : tensor<f32> loc(#loc5310)
    %25912 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25913 = stablehlo.broadcast_in_dim %25912, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25914 = stablehlo.broadcast_in_dim %25911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25915 = stablehlo.concatenate %25913, %25914, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25916 = stablehlo.abs %25418 : tensor<1x2048xf32> loc(#loc5307)
    %25917 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25918 = stablehlo.reduce(%25916 init: %25917) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25919 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25920 = stablehlo.maximum %25918, %25919 : tensor<f32> loc(#loc5309)
    %25921 = stablehlo.broadcast_in_dim %25920, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25922 = stablehlo.divide %25418, %25921 : tensor<1x2048xf32> loc(#loc5314)
    %25923 = stablehlo.multiply %25920, %25920 : tensor<f32> loc(#loc5310)
    %25924 = stablehlo.multiply %25922, %25922 : tensor<1x2048xf32> loc(#loc5310)
    %25925 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25926 = stablehlo.reduce(%25924 init: %25925) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25927 = stablehlo.multiply %25923, %25926 : tensor<f32> loc(#loc5310)
    %25928 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25929 = stablehlo.broadcast_in_dim %25928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25930 = stablehlo.broadcast_in_dim %25927, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25931 = stablehlo.concatenate %25929, %25930, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25932 = stablehlo.abs %25419 : tensor<1x2048xf32> loc(#loc5307)
    %25933 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25934 = stablehlo.reduce(%25932 init: %25933) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25935 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25936 = stablehlo.maximum %25934, %25935 : tensor<f32> loc(#loc5309)
    %25937 = stablehlo.broadcast_in_dim %25936, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25938 = stablehlo.divide %25419, %25937 : tensor<1x2048xf32> loc(#loc5314)
    %25939 = stablehlo.multiply %25936, %25936 : tensor<f32> loc(#loc5310)
    %25940 = stablehlo.multiply %25938, %25938 : tensor<1x2048xf32> loc(#loc5310)
    %25941 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25942 = stablehlo.reduce(%25940 init: %25941) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25943 = stablehlo.multiply %25939, %25942 : tensor<f32> loc(#loc5310)
    %25944 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25945 = stablehlo.broadcast_in_dim %25944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25946 = stablehlo.broadcast_in_dim %25943, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25947 = stablehlo.concatenate %25945, %25946, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25948 = stablehlo.abs %25420 : tensor<1x2048xf32> loc(#loc5307)
    %25949 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25950 = stablehlo.reduce(%25948 init: %25949) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25951 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25952 = stablehlo.maximum %25950, %25951 : tensor<f32> loc(#loc5309)
    %25953 = stablehlo.broadcast_in_dim %25952, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %25954 = stablehlo.divide %25420, %25953 : tensor<1x2048xf32> loc(#loc5314)
    %25955 = stablehlo.multiply %25952, %25952 : tensor<f32> loc(#loc5310)
    %25956 = stablehlo.multiply %25954, %25954 : tensor<1x2048xf32> loc(#loc5310)
    %25957 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %25958 = stablehlo.reduce(%25956 init: %25957) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %25959 = stablehlo.multiply %25955, %25958 : tensor<f32> loc(#loc5310)
    %25960 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %25961 = stablehlo.broadcast_in_dim %25960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25962 = stablehlo.broadcast_in_dim %25959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25963 = stablehlo.concatenate %25961, %25962, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25964 = stablehlo.abs %25421 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %25965 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %25966 = stablehlo.reduce(%25964 init: %25965) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %25967 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25968 = stablehlo.maximum %25966, %25967 : tensor<f32> loc(#loc5309)
    %25969 = stablehlo.broadcast_in_dim %25968, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %25970 = stablehlo.divide %25421, %25969 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %25971 = stablehlo.multiply %25968, %25968 : tensor<f32> loc(#loc5310)
    %25972 = stablehlo.multiply %25970, %25970 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %25973 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %25974 = stablehlo.reduce(%25972 init: %25973) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %25975 = stablehlo.multiply %25971, %25974 : tensor<f32> loc(#loc5310)
    %25976 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %25977 = stablehlo.broadcast_in_dim %25976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25978 = stablehlo.broadcast_in_dim %25975, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25979 = stablehlo.concatenate %25977, %25978, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25980 = stablehlo.abs %25422 : tensor<1x64xf32> loc(#loc5307)
    %25981 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %25982 = stablehlo.reduce(%25980 init: %25981) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %25983 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %25984 = stablehlo.maximum %25982, %25983 : tensor<f32> loc(#loc5309)
    %25985 = stablehlo.broadcast_in_dim %25984, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %25986 = stablehlo.divide %25422, %25985 : tensor<1x64xf32> loc(#loc5314)
    %25987 = stablehlo.multiply %25984, %25984 : tensor<f32> loc(#loc5310)
    %25988 = stablehlo.multiply %25986, %25986 : tensor<1x64xf32> loc(#loc5310)
    %25989 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %25990 = stablehlo.reduce(%25988 init: %25989) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %25991 = stablehlo.multiply %25987, %25990 : tensor<f32> loc(#loc5310)
    %25992 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %25993 = stablehlo.broadcast_in_dim %25992, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25994 = stablehlo.broadcast_in_dim %25991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %25995 = stablehlo.concatenate %25993, %25994, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %25996 = stablehlo.abs %25423 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %25997 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %25998 = stablehlo.reduce(%25996 init: %25997) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %25999 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26000 = stablehlo.maximum %25998, %25999 : tensor<f32> loc(#loc5309)
    %26001 = stablehlo.broadcast_in_dim %26000, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %26002 = stablehlo.divide %25423, %26001 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %26003 = stablehlo.multiply %26000, %26000 : tensor<f32> loc(#loc5310)
    %26004 = stablehlo.multiply %26002, %26002 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %26005 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %26006 = stablehlo.reduce(%26004 init: %26005) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %26007 = stablehlo.multiply %26003, %26006 : tensor<f32> loc(#loc5310)
    %26008 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %26009 = stablehlo.broadcast_in_dim %26008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26010 = stablehlo.broadcast_in_dim %26007, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26011 = stablehlo.concatenate %26009, %26010, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26012 = stablehlo.abs %25424 : tensor<1x8192xf32> loc(#loc5307)
    %26013 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26014 = stablehlo.reduce(%26012 init: %26013) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26015 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26016 = stablehlo.maximum %26014, %26015 : tensor<f32> loc(#loc5309)
    %26017 = stablehlo.broadcast_in_dim %26016, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %26018 = stablehlo.divide %25424, %26017 : tensor<1x8192xf32> loc(#loc5314)
    %26019 = stablehlo.multiply %26016, %26016 : tensor<f32> loc(#loc5310)
    %26020 = stablehlo.multiply %26018, %26018 : tensor<1x8192xf32> loc(#loc5310)
    %26021 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26022 = stablehlo.reduce(%26020 init: %26021) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26023 = stablehlo.multiply %26019, %26022 : tensor<f32> loc(#loc5310)
    %26024 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %26025 = stablehlo.broadcast_in_dim %26024, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26026 = stablehlo.broadcast_in_dim %26023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26027 = stablehlo.concatenate %26025, %26026, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26028 = stablehlo.abs %25425 : tensor<1x2048x8192xf32> loc(#loc5307)
    %26029 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26030 = stablehlo.reduce(%26028 init: %26029) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26032 = stablehlo.maximum %26030, %26031 : tensor<f32> loc(#loc5309)
    %26033 = stablehlo.broadcast_in_dim %26032, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %26034 = stablehlo.divide %25425, %26033 : tensor<1x2048x8192xf32> loc(#loc5314)
    %26035 = stablehlo.multiply %26032, %26032 : tensor<f32> loc(#loc5310)
    %26036 = stablehlo.multiply %26034, %26034 : tensor<1x2048x8192xf32> loc(#loc5310)
    %26037 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26038 = stablehlo.reduce(%26036 init: %26037) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26039 = stablehlo.multiply %26035, %26038 : tensor<f32> loc(#loc5310)
    %26040 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26041 = stablehlo.broadcast_in_dim %26040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26042 = stablehlo.broadcast_in_dim %26039, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26043 = stablehlo.concatenate %26041, %26042, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26044 = stablehlo.abs %25426 : tensor<1x2048xf32> loc(#loc5307)
    %26045 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26046 = stablehlo.reduce(%26044 init: %26045) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26047 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26048 = stablehlo.maximum %26046, %26047 : tensor<f32> loc(#loc5309)
    %26049 = stablehlo.broadcast_in_dim %26048, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26050 = stablehlo.divide %25426, %26049 : tensor<1x2048xf32> loc(#loc5314)
    %26051 = stablehlo.multiply %26048, %26048 : tensor<f32> loc(#loc5310)
    %26052 = stablehlo.multiply %26050, %26050 : tensor<1x2048xf32> loc(#loc5310)
    %26053 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26054 = stablehlo.reduce(%26052 init: %26053) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26055 = stablehlo.multiply %26051, %26054 : tensor<f32> loc(#loc5310)
    %26056 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26057 = stablehlo.broadcast_in_dim %26056, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26058 = stablehlo.broadcast_in_dim %26055, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26059 = stablehlo.concatenate %26057, %26058, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26060 = stablehlo.abs %25427 : tensor<1x8192x2048xf32> loc(#loc5307)
    %26061 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26062 = stablehlo.reduce(%26060 init: %26061) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26064 = stablehlo.maximum %26062, %26063 : tensor<f32> loc(#loc5309)
    %26065 = stablehlo.broadcast_in_dim %26064, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %26066 = stablehlo.divide %25427, %26065 : tensor<1x8192x2048xf32> loc(#loc5314)
    %26067 = stablehlo.multiply %26064, %26064 : tensor<f32> loc(#loc5310)
    %26068 = stablehlo.multiply %26066, %26066 : tensor<1x8192x2048xf32> loc(#loc5310)
    %26069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26070 = stablehlo.reduce(%26068 init: %26069) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26071 = stablehlo.multiply %26067, %26070 : tensor<f32> loc(#loc5310)
    %26072 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26073 = stablehlo.broadcast_in_dim %26072, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26074 = stablehlo.broadcast_in_dim %26071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26075 = stablehlo.concatenate %26073, %26074, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26076 = stablehlo.abs %25428 : tensor<1x2048xf32> loc(#loc5307)
    %26077 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26078 = stablehlo.reduce(%26076 init: %26077) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26079 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26080 = stablehlo.maximum %26078, %26079 : tensor<f32> loc(#loc5309)
    %26081 = stablehlo.broadcast_in_dim %26080, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26082 = stablehlo.divide %25428, %26081 : tensor<1x2048xf32> loc(#loc5314)
    %26083 = stablehlo.multiply %26080, %26080 : tensor<f32> loc(#loc5310)
    %26084 = stablehlo.multiply %26082, %26082 : tensor<1x2048xf32> loc(#loc5310)
    %26085 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26086 = stablehlo.reduce(%26084 init: %26085) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26087 = stablehlo.multiply %26083, %26086 : tensor<f32> loc(#loc5310)
    %26088 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26089 = stablehlo.broadcast_in_dim %26088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26090 = stablehlo.broadcast_in_dim %26087, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26091 = stablehlo.concatenate %26089, %26090, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26092 = stablehlo.abs %25429 : tensor<1x2048xf32> loc(#loc5307)
    %26093 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26094 = stablehlo.reduce(%26092 init: %26093) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26095 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26096 = stablehlo.maximum %26094, %26095 : tensor<f32> loc(#loc5309)
    %26097 = stablehlo.broadcast_in_dim %26096, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26098 = stablehlo.divide %25429, %26097 : tensor<1x2048xf32> loc(#loc5314)
    %26099 = stablehlo.multiply %26096, %26096 : tensor<f32> loc(#loc5310)
    %26100 = stablehlo.multiply %26098, %26098 : tensor<1x2048xf32> loc(#loc5310)
    %26101 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26102 = stablehlo.reduce(%26100 init: %26101) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26103 = stablehlo.multiply %26099, %26102 : tensor<f32> loc(#loc5310)
    %26104 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26105 = stablehlo.broadcast_in_dim %26104, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26106 = stablehlo.broadcast_in_dim %26103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26107 = stablehlo.concatenate %26105, %26106, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26108 = stablehlo.abs %25430 : tensor<1x2048xf32> loc(#loc5307)
    %26109 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26110 = stablehlo.reduce(%26108 init: %26109) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26112 = stablehlo.maximum %26110, %26111 : tensor<f32> loc(#loc5309)
    %26113 = stablehlo.broadcast_in_dim %26112, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26114 = stablehlo.divide %25430, %26113 : tensor<1x2048xf32> loc(#loc5314)
    %26115 = stablehlo.multiply %26112, %26112 : tensor<f32> loc(#loc5310)
    %26116 = stablehlo.multiply %26114, %26114 : tensor<1x2048xf32> loc(#loc5310)
    %26117 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26118 = stablehlo.reduce(%26116 init: %26117) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26119 = stablehlo.multiply %26115, %26118 : tensor<f32> loc(#loc5310)
    %26120 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26121 = stablehlo.broadcast_in_dim %26120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26122 = stablehlo.broadcast_in_dim %26119, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26123 = stablehlo.concatenate %26121, %26122, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26124 = stablehlo.abs %25431 : tensor<1x2048xf32> loc(#loc5307)
    %26125 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26126 = stablehlo.reduce(%26124 init: %26125) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26127 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26128 = stablehlo.maximum %26126, %26127 : tensor<f32> loc(#loc5309)
    %26129 = stablehlo.broadcast_in_dim %26128, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26130 = stablehlo.divide %25431, %26129 : tensor<1x2048xf32> loc(#loc5314)
    %26131 = stablehlo.multiply %26128, %26128 : tensor<f32> loc(#loc5310)
    %26132 = stablehlo.multiply %26130, %26130 : tensor<1x2048xf32> loc(#loc5310)
    %26133 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26134 = stablehlo.reduce(%26132 init: %26133) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26135 = stablehlo.multiply %26131, %26134 : tensor<f32> loc(#loc5310)
    %26136 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26137 = stablehlo.broadcast_in_dim %26136, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26138 = stablehlo.broadcast_in_dim %26135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26139 = stablehlo.concatenate %26137, %26138, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26140 = stablehlo.abs %25432 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %26141 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %26142 = stablehlo.reduce(%26140 init: %26141) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %26143 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26144 = stablehlo.maximum %26142, %26143 : tensor<f32> loc(#loc5309)
    %26145 = stablehlo.broadcast_in_dim %26144, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26146 = stablehlo.divide %25432, %26145 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26147 = stablehlo.multiply %26144, %26144 : tensor<f32> loc(#loc5310)
    %26148 = stablehlo.multiply %26146, %26146 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %26149 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %26150 = stablehlo.reduce(%26148 init: %26149) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %26151 = stablehlo.multiply %26147, %26150 : tensor<f32> loc(#loc5310)
    %26152 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %26153 = stablehlo.broadcast_in_dim %26152, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26154 = stablehlo.broadcast_in_dim %26151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26155 = stablehlo.concatenate %26153, %26154, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26156 = stablehlo.abs %25433 : tensor<1x64xf32> loc(#loc5307)
    %26157 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26158 = stablehlo.reduce(%26156 init: %26157) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26159 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26160 = stablehlo.maximum %26158, %26159 : tensor<f32> loc(#loc5309)
    %26161 = stablehlo.broadcast_in_dim %26160, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %26162 = stablehlo.divide %25433, %26161 : tensor<1x64xf32> loc(#loc5314)
    %26163 = stablehlo.multiply %26160, %26160 : tensor<f32> loc(#loc5310)
    %26164 = stablehlo.multiply %26162, %26162 : tensor<1x64xf32> loc(#loc5310)
    %26165 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26166 = stablehlo.reduce(%26164 init: %26165) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26167 = stablehlo.multiply %26163, %26166 : tensor<f32> loc(#loc5310)
    %26168 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %26169 = stablehlo.broadcast_in_dim %26168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26170 = stablehlo.broadcast_in_dim %26167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26171 = stablehlo.concatenate %26169, %26170, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26172 = stablehlo.abs %25434 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %26173 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %26174 = stablehlo.reduce(%26172 init: %26173) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %26175 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26176 = stablehlo.maximum %26174, %26175 : tensor<f32> loc(#loc5309)
    %26177 = stablehlo.broadcast_in_dim %26176, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %26178 = stablehlo.divide %25434, %26177 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %26179 = stablehlo.multiply %26176, %26176 : tensor<f32> loc(#loc5310)
    %26180 = stablehlo.multiply %26178, %26178 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %26181 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %26182 = stablehlo.reduce(%26180 init: %26181) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %26183 = stablehlo.multiply %26179, %26182 : tensor<f32> loc(#loc5310)
    %26184 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %26185 = stablehlo.broadcast_in_dim %26184, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26186 = stablehlo.broadcast_in_dim %26183, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26187 = stablehlo.concatenate %26185, %26186, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26188 = stablehlo.abs %25435 : tensor<1x8192xf32> loc(#loc5307)
    %26189 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26190 = stablehlo.reduce(%26188 init: %26189) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26191 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26192 = stablehlo.maximum %26190, %26191 : tensor<f32> loc(#loc5309)
    %26193 = stablehlo.broadcast_in_dim %26192, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %26194 = stablehlo.divide %25435, %26193 : tensor<1x8192xf32> loc(#loc5314)
    %26195 = stablehlo.multiply %26192, %26192 : tensor<f32> loc(#loc5310)
    %26196 = stablehlo.multiply %26194, %26194 : tensor<1x8192xf32> loc(#loc5310)
    %26197 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26198 = stablehlo.reduce(%26196 init: %26197) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26199 = stablehlo.multiply %26195, %26198 : tensor<f32> loc(#loc5310)
    %26200 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %26201 = stablehlo.broadcast_in_dim %26200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26202 = stablehlo.broadcast_in_dim %26199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26203 = stablehlo.concatenate %26201, %26202, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26204 = stablehlo.abs %25436 : tensor<1x2048x8192xf32> loc(#loc5307)
    %26205 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26206 = stablehlo.reduce(%26204 init: %26205) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26208 = stablehlo.maximum %26206, %26207 : tensor<f32> loc(#loc5309)
    %26209 = stablehlo.broadcast_in_dim %26208, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %26210 = stablehlo.divide %25436, %26209 : tensor<1x2048x8192xf32> loc(#loc5314)
    %26211 = stablehlo.multiply %26208, %26208 : tensor<f32> loc(#loc5310)
    %26212 = stablehlo.multiply %26210, %26210 : tensor<1x2048x8192xf32> loc(#loc5310)
    %26213 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26214 = stablehlo.reduce(%26212 init: %26213) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26215 = stablehlo.multiply %26211, %26214 : tensor<f32> loc(#loc5310)
    %26216 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26217 = stablehlo.broadcast_in_dim %26216, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26218 = stablehlo.broadcast_in_dim %26215, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26219 = stablehlo.concatenate %26217, %26218, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26220 = stablehlo.abs %25437 : tensor<1x2048xf32> loc(#loc5307)
    %26221 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26222 = stablehlo.reduce(%26220 init: %26221) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26223 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26224 = stablehlo.maximum %26222, %26223 : tensor<f32> loc(#loc5309)
    %26225 = stablehlo.broadcast_in_dim %26224, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26226 = stablehlo.divide %25437, %26225 : tensor<1x2048xf32> loc(#loc5314)
    %26227 = stablehlo.multiply %26224, %26224 : tensor<f32> loc(#loc5310)
    %26228 = stablehlo.multiply %26226, %26226 : tensor<1x2048xf32> loc(#loc5310)
    %26229 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26230 = stablehlo.reduce(%26228 init: %26229) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26231 = stablehlo.multiply %26227, %26230 : tensor<f32> loc(#loc5310)
    %26232 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26233 = stablehlo.broadcast_in_dim %26232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26234 = stablehlo.broadcast_in_dim %26231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26235 = stablehlo.concatenate %26233, %26234, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26236 = stablehlo.abs %25438 : tensor<1x8192x2048xf32> loc(#loc5307)
    %26237 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26238 = stablehlo.reduce(%26236 init: %26237) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26239 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26240 = stablehlo.maximum %26238, %26239 : tensor<f32> loc(#loc5309)
    %26241 = stablehlo.broadcast_in_dim %26240, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %26242 = stablehlo.divide %25438, %26241 : tensor<1x8192x2048xf32> loc(#loc5314)
    %26243 = stablehlo.multiply %26240, %26240 : tensor<f32> loc(#loc5310)
    %26244 = stablehlo.multiply %26242, %26242 : tensor<1x8192x2048xf32> loc(#loc5310)
    %26245 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26246 = stablehlo.reduce(%26244 init: %26245) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26247 = stablehlo.multiply %26243, %26246 : tensor<f32> loc(#loc5310)
    %26248 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26249 = stablehlo.broadcast_in_dim %26248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26250 = stablehlo.broadcast_in_dim %26247, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26251 = stablehlo.concatenate %26249, %26250, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26252 = stablehlo.abs %25439 : tensor<1x2048xf32> loc(#loc5307)
    %26253 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26254 = stablehlo.reduce(%26252 init: %26253) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26256 = stablehlo.maximum %26254, %26255 : tensor<f32> loc(#loc5309)
    %26257 = stablehlo.broadcast_in_dim %26256, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26258 = stablehlo.divide %25439, %26257 : tensor<1x2048xf32> loc(#loc5314)
    %26259 = stablehlo.multiply %26256, %26256 : tensor<f32> loc(#loc5310)
    %26260 = stablehlo.multiply %26258, %26258 : tensor<1x2048xf32> loc(#loc5310)
    %26261 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26262 = stablehlo.reduce(%26260 init: %26261) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26263 = stablehlo.multiply %26259, %26262 : tensor<f32> loc(#loc5310)
    %26264 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26265 = stablehlo.broadcast_in_dim %26264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26266 = stablehlo.broadcast_in_dim %26263, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26267 = stablehlo.concatenate %26265, %26266, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26268 = stablehlo.abs %25440 : tensor<1x2048xf32> loc(#loc5307)
    %26269 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26270 = stablehlo.reduce(%26268 init: %26269) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26272 = stablehlo.maximum %26270, %26271 : tensor<f32> loc(#loc5309)
    %26273 = stablehlo.broadcast_in_dim %26272, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26274 = stablehlo.divide %25440, %26273 : tensor<1x2048xf32> loc(#loc5314)
    %26275 = stablehlo.multiply %26272, %26272 : tensor<f32> loc(#loc5310)
    %26276 = stablehlo.multiply %26274, %26274 : tensor<1x2048xf32> loc(#loc5310)
    %26277 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26278 = stablehlo.reduce(%26276 init: %26277) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26279 = stablehlo.multiply %26275, %26278 : tensor<f32> loc(#loc5310)
    %26280 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26281 = stablehlo.broadcast_in_dim %26280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26282 = stablehlo.broadcast_in_dim %26279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26283 = stablehlo.concatenate %26281, %26282, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26284 = stablehlo.abs %25441 : tensor<1x2048xf32> loc(#loc5307)
    %26285 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26286 = stablehlo.reduce(%26284 init: %26285) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26287 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26288 = stablehlo.maximum %26286, %26287 : tensor<f32> loc(#loc5309)
    %26289 = stablehlo.broadcast_in_dim %26288, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26290 = stablehlo.divide %25441, %26289 : tensor<1x2048xf32> loc(#loc5314)
    %26291 = stablehlo.multiply %26288, %26288 : tensor<f32> loc(#loc5310)
    %26292 = stablehlo.multiply %26290, %26290 : tensor<1x2048xf32> loc(#loc5310)
    %26293 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26294 = stablehlo.reduce(%26292 init: %26293) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26295 = stablehlo.multiply %26291, %26294 : tensor<f32> loc(#loc5310)
    %26296 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26297 = stablehlo.broadcast_in_dim %26296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26298 = stablehlo.broadcast_in_dim %26295, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26299 = stablehlo.concatenate %26297, %26298, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26300 = stablehlo.abs %25442 : tensor<1x2048xf32> loc(#loc5307)
    %26301 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26302 = stablehlo.reduce(%26300 init: %26301) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26303 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26304 = stablehlo.maximum %26302, %26303 : tensor<f32> loc(#loc5309)
    %26305 = stablehlo.broadcast_in_dim %26304, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26306 = stablehlo.divide %25442, %26305 : tensor<1x2048xf32> loc(#loc5314)
    %26307 = stablehlo.multiply %26304, %26304 : tensor<f32> loc(#loc5310)
    %26308 = stablehlo.multiply %26306, %26306 : tensor<1x2048xf32> loc(#loc5310)
    %26309 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26310 = stablehlo.reduce(%26308 init: %26309) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26311 = stablehlo.multiply %26307, %26310 : tensor<f32> loc(#loc5310)
    %26312 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26313 = stablehlo.broadcast_in_dim %26312, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26314 = stablehlo.broadcast_in_dim %26311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26315 = stablehlo.concatenate %26313, %26314, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26316 = stablehlo.abs %25443 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %26317 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %26318 = stablehlo.reduce(%26316 init: %26317) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %26319 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26320 = stablehlo.maximum %26318, %26319 : tensor<f32> loc(#loc5309)
    %26321 = stablehlo.broadcast_in_dim %26320, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26322 = stablehlo.divide %25443, %26321 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26323 = stablehlo.multiply %26320, %26320 : tensor<f32> loc(#loc5310)
    %26324 = stablehlo.multiply %26322, %26322 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %26325 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %26326 = stablehlo.reduce(%26324 init: %26325) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %26327 = stablehlo.multiply %26323, %26326 : tensor<f32> loc(#loc5310)
    %26328 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %26329 = stablehlo.broadcast_in_dim %26328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26330 = stablehlo.broadcast_in_dim %26327, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26331 = stablehlo.concatenate %26329, %26330, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26332 = stablehlo.abs %25444 : tensor<1x64xf32> loc(#loc5307)
    %26333 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26334 = stablehlo.reduce(%26332 init: %26333) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26335 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26336 = stablehlo.maximum %26334, %26335 : tensor<f32> loc(#loc5309)
    %26337 = stablehlo.broadcast_in_dim %26336, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %26338 = stablehlo.divide %25444, %26337 : tensor<1x64xf32> loc(#loc5314)
    %26339 = stablehlo.multiply %26336, %26336 : tensor<f32> loc(#loc5310)
    %26340 = stablehlo.multiply %26338, %26338 : tensor<1x64xf32> loc(#loc5310)
    %26341 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26342 = stablehlo.reduce(%26340 init: %26341) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26343 = stablehlo.multiply %26339, %26342 : tensor<f32> loc(#loc5310)
    %26344 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %26345 = stablehlo.broadcast_in_dim %26344, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26346 = stablehlo.broadcast_in_dim %26343, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26347 = stablehlo.concatenate %26345, %26346, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26348 = stablehlo.abs %25445 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %26349 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %26350 = stablehlo.reduce(%26348 init: %26349) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %26351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26352 = stablehlo.maximum %26350, %26351 : tensor<f32> loc(#loc5309)
    %26353 = stablehlo.broadcast_in_dim %26352, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %26354 = stablehlo.divide %25445, %26353 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %26355 = stablehlo.multiply %26352, %26352 : tensor<f32> loc(#loc5310)
    %26356 = stablehlo.multiply %26354, %26354 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %26357 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %26358 = stablehlo.reduce(%26356 init: %26357) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %26359 = stablehlo.multiply %26355, %26358 : tensor<f32> loc(#loc5310)
    %26360 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %26361 = stablehlo.broadcast_in_dim %26360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26362 = stablehlo.broadcast_in_dim %26359, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26363 = stablehlo.concatenate %26361, %26362, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26364 = stablehlo.abs %25446 : tensor<1x8192xf32> loc(#loc5307)
    %26365 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26366 = stablehlo.reduce(%26364 init: %26365) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26367 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26368 = stablehlo.maximum %26366, %26367 : tensor<f32> loc(#loc5309)
    %26369 = stablehlo.broadcast_in_dim %26368, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %26370 = stablehlo.divide %25446, %26369 : tensor<1x8192xf32> loc(#loc5314)
    %26371 = stablehlo.multiply %26368, %26368 : tensor<f32> loc(#loc5310)
    %26372 = stablehlo.multiply %26370, %26370 : tensor<1x8192xf32> loc(#loc5310)
    %26373 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26374 = stablehlo.reduce(%26372 init: %26373) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26375 = stablehlo.multiply %26371, %26374 : tensor<f32> loc(#loc5310)
    %26376 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %26377 = stablehlo.broadcast_in_dim %26376, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26378 = stablehlo.broadcast_in_dim %26375, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26379 = stablehlo.concatenate %26377, %26378, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26380 = stablehlo.abs %25447 : tensor<1x2048x8192xf32> loc(#loc5307)
    %26381 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26382 = stablehlo.reduce(%26380 init: %26381) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26383 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26384 = stablehlo.maximum %26382, %26383 : tensor<f32> loc(#loc5309)
    %26385 = stablehlo.broadcast_in_dim %26384, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %26386 = stablehlo.divide %25447, %26385 : tensor<1x2048x8192xf32> loc(#loc5314)
    %26387 = stablehlo.multiply %26384, %26384 : tensor<f32> loc(#loc5310)
    %26388 = stablehlo.multiply %26386, %26386 : tensor<1x2048x8192xf32> loc(#loc5310)
    %26389 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26390 = stablehlo.reduce(%26388 init: %26389) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26391 = stablehlo.multiply %26387, %26390 : tensor<f32> loc(#loc5310)
    %26392 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26393 = stablehlo.broadcast_in_dim %26392, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26394 = stablehlo.broadcast_in_dim %26391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26395 = stablehlo.concatenate %26393, %26394, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26396 = stablehlo.abs %25448 : tensor<1x2048xf32> loc(#loc5307)
    %26397 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26398 = stablehlo.reduce(%26396 init: %26397) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26399 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26400 = stablehlo.maximum %26398, %26399 : tensor<f32> loc(#loc5309)
    %26401 = stablehlo.broadcast_in_dim %26400, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26402 = stablehlo.divide %25448, %26401 : tensor<1x2048xf32> loc(#loc5314)
    %26403 = stablehlo.multiply %26400, %26400 : tensor<f32> loc(#loc5310)
    %26404 = stablehlo.multiply %26402, %26402 : tensor<1x2048xf32> loc(#loc5310)
    %26405 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26406 = stablehlo.reduce(%26404 init: %26405) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26407 = stablehlo.multiply %26403, %26406 : tensor<f32> loc(#loc5310)
    %26408 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26409 = stablehlo.broadcast_in_dim %26408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26410 = stablehlo.broadcast_in_dim %26407, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26411 = stablehlo.concatenate %26409, %26410, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26412 = stablehlo.abs %25449 : tensor<1x8192x2048xf32> loc(#loc5307)
    %26413 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26414 = stablehlo.reduce(%26412 init: %26413) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26415 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26416 = stablehlo.maximum %26414, %26415 : tensor<f32> loc(#loc5309)
    %26417 = stablehlo.broadcast_in_dim %26416, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %26418 = stablehlo.divide %25449, %26417 : tensor<1x8192x2048xf32> loc(#loc5314)
    %26419 = stablehlo.multiply %26416, %26416 : tensor<f32> loc(#loc5310)
    %26420 = stablehlo.multiply %26418, %26418 : tensor<1x8192x2048xf32> loc(#loc5310)
    %26421 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26422 = stablehlo.reduce(%26420 init: %26421) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26423 = stablehlo.multiply %26419, %26422 : tensor<f32> loc(#loc5310)
    %26424 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26425 = stablehlo.broadcast_in_dim %26424, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26426 = stablehlo.broadcast_in_dim %26423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26427 = stablehlo.concatenate %26425, %26426, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26428 = stablehlo.abs %25450 : tensor<1x2048xf32> loc(#loc5307)
    %26429 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26430 = stablehlo.reduce(%26428 init: %26429) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26432 = stablehlo.maximum %26430, %26431 : tensor<f32> loc(#loc5309)
    %26433 = stablehlo.broadcast_in_dim %26432, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26434 = stablehlo.divide %25450, %26433 : tensor<1x2048xf32> loc(#loc5314)
    %26435 = stablehlo.multiply %26432, %26432 : tensor<f32> loc(#loc5310)
    %26436 = stablehlo.multiply %26434, %26434 : tensor<1x2048xf32> loc(#loc5310)
    %26437 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26438 = stablehlo.reduce(%26436 init: %26437) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26439 = stablehlo.multiply %26435, %26438 : tensor<f32> loc(#loc5310)
    %26440 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26441 = stablehlo.broadcast_in_dim %26440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26442 = stablehlo.broadcast_in_dim %26439, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26443 = stablehlo.concatenate %26441, %26442, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26444 = stablehlo.abs %25451 : tensor<1x2048xf32> loc(#loc5307)
    %26445 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26446 = stablehlo.reduce(%26444 init: %26445) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26447 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26448 = stablehlo.maximum %26446, %26447 : tensor<f32> loc(#loc5309)
    %26449 = stablehlo.broadcast_in_dim %26448, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26450 = stablehlo.divide %25451, %26449 : tensor<1x2048xf32> loc(#loc5314)
    %26451 = stablehlo.multiply %26448, %26448 : tensor<f32> loc(#loc5310)
    %26452 = stablehlo.multiply %26450, %26450 : tensor<1x2048xf32> loc(#loc5310)
    %26453 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26454 = stablehlo.reduce(%26452 init: %26453) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26455 = stablehlo.multiply %26451, %26454 : tensor<f32> loc(#loc5310)
    %26456 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26457 = stablehlo.broadcast_in_dim %26456, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26458 = stablehlo.broadcast_in_dim %26455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26459 = stablehlo.concatenate %26457, %26458, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26460 = stablehlo.abs %25452 : tensor<1x2048xf32> loc(#loc5307)
    %26461 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26462 = stablehlo.reduce(%26460 init: %26461) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26463 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26464 = stablehlo.maximum %26462, %26463 : tensor<f32> loc(#loc5309)
    %26465 = stablehlo.broadcast_in_dim %26464, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26466 = stablehlo.divide %25452, %26465 : tensor<1x2048xf32> loc(#loc5314)
    %26467 = stablehlo.multiply %26464, %26464 : tensor<f32> loc(#loc5310)
    %26468 = stablehlo.multiply %26466, %26466 : tensor<1x2048xf32> loc(#loc5310)
    %26469 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26470 = stablehlo.reduce(%26468 init: %26469) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26471 = stablehlo.multiply %26467, %26470 : tensor<f32> loc(#loc5310)
    %26472 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26473 = stablehlo.broadcast_in_dim %26472, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26474 = stablehlo.broadcast_in_dim %26471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26475 = stablehlo.concatenate %26473, %26474, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26476 = stablehlo.abs %25453 : tensor<1x2048xf32> loc(#loc5307)
    %26477 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26478 = stablehlo.reduce(%26476 init: %26477) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26479 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26480 = stablehlo.maximum %26478, %26479 : tensor<f32> loc(#loc5309)
    %26481 = stablehlo.broadcast_in_dim %26480, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26482 = stablehlo.divide %25453, %26481 : tensor<1x2048xf32> loc(#loc5314)
    %26483 = stablehlo.multiply %26480, %26480 : tensor<f32> loc(#loc5310)
    %26484 = stablehlo.multiply %26482, %26482 : tensor<1x2048xf32> loc(#loc5310)
    %26485 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26486 = stablehlo.reduce(%26484 init: %26485) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26487 = stablehlo.multiply %26483, %26486 : tensor<f32> loc(#loc5310)
    %26488 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26489 = stablehlo.broadcast_in_dim %26488, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26490 = stablehlo.broadcast_in_dim %26487, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26491 = stablehlo.concatenate %26489, %26490, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26492 = stablehlo.abs %25454 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %26493 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %26494 = stablehlo.reduce(%26492 init: %26493) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %26495 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26496 = stablehlo.maximum %26494, %26495 : tensor<f32> loc(#loc5309)
    %26497 = stablehlo.broadcast_in_dim %26496, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26498 = stablehlo.divide %25454, %26497 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26499 = stablehlo.multiply %26496, %26496 : tensor<f32> loc(#loc5310)
    %26500 = stablehlo.multiply %26498, %26498 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %26501 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %26502 = stablehlo.reduce(%26500 init: %26501) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %26503 = stablehlo.multiply %26499, %26502 : tensor<f32> loc(#loc5310)
    %26504 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %26505 = stablehlo.broadcast_in_dim %26504, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26506 = stablehlo.broadcast_in_dim %26503, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26507 = stablehlo.concatenate %26505, %26506, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26508 = stablehlo.abs %25455 : tensor<1x64xf32> loc(#loc5307)
    %26509 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26510 = stablehlo.reduce(%26508 init: %26509) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26511 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26512 = stablehlo.maximum %26510, %26511 : tensor<f32> loc(#loc5309)
    %26513 = stablehlo.broadcast_in_dim %26512, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %26514 = stablehlo.divide %25455, %26513 : tensor<1x64xf32> loc(#loc5314)
    %26515 = stablehlo.multiply %26512, %26512 : tensor<f32> loc(#loc5310)
    %26516 = stablehlo.multiply %26514, %26514 : tensor<1x64xf32> loc(#loc5310)
    %26517 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26518 = stablehlo.reduce(%26516 init: %26517) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26519 = stablehlo.multiply %26515, %26518 : tensor<f32> loc(#loc5310)
    %26520 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %26521 = stablehlo.broadcast_in_dim %26520, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26522 = stablehlo.broadcast_in_dim %26519, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26523 = stablehlo.concatenate %26521, %26522, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26524 = stablehlo.abs %25456 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %26525 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %26526 = stablehlo.reduce(%26524 init: %26525) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %26527 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26528 = stablehlo.maximum %26526, %26527 : tensor<f32> loc(#loc5309)
    %26529 = stablehlo.broadcast_in_dim %26528, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %26530 = stablehlo.divide %25456, %26529 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %26531 = stablehlo.multiply %26528, %26528 : tensor<f32> loc(#loc5310)
    %26532 = stablehlo.multiply %26530, %26530 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %26533 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %26534 = stablehlo.reduce(%26532 init: %26533) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %26535 = stablehlo.multiply %26531, %26534 : tensor<f32> loc(#loc5310)
    %26536 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %26537 = stablehlo.broadcast_in_dim %26536, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26538 = stablehlo.broadcast_in_dim %26535, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26539 = stablehlo.concatenate %26537, %26538, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26540 = stablehlo.abs %25457 : tensor<1x8192xf32> loc(#loc5307)
    %26541 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26542 = stablehlo.reduce(%26540 init: %26541) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26543 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26544 = stablehlo.maximum %26542, %26543 : tensor<f32> loc(#loc5309)
    %26545 = stablehlo.broadcast_in_dim %26544, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %26546 = stablehlo.divide %25457, %26545 : tensor<1x8192xf32> loc(#loc5314)
    %26547 = stablehlo.multiply %26544, %26544 : tensor<f32> loc(#loc5310)
    %26548 = stablehlo.multiply %26546, %26546 : tensor<1x8192xf32> loc(#loc5310)
    %26549 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26550 = stablehlo.reduce(%26548 init: %26549) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26551 = stablehlo.multiply %26547, %26550 : tensor<f32> loc(#loc5310)
    %26552 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %26553 = stablehlo.broadcast_in_dim %26552, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26554 = stablehlo.broadcast_in_dim %26551, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26555 = stablehlo.concatenate %26553, %26554, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26556 = stablehlo.abs %25458 : tensor<1x2048x8192xf32> loc(#loc5307)
    %26557 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26558 = stablehlo.reduce(%26556 init: %26557) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26559 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26560 = stablehlo.maximum %26558, %26559 : tensor<f32> loc(#loc5309)
    %26561 = stablehlo.broadcast_in_dim %26560, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %26562 = stablehlo.divide %25458, %26561 : tensor<1x2048x8192xf32> loc(#loc5314)
    %26563 = stablehlo.multiply %26560, %26560 : tensor<f32> loc(#loc5310)
    %26564 = stablehlo.multiply %26562, %26562 : tensor<1x2048x8192xf32> loc(#loc5310)
    %26565 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26566 = stablehlo.reduce(%26564 init: %26565) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26567 = stablehlo.multiply %26563, %26566 : tensor<f32> loc(#loc5310)
    %26568 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26569 = stablehlo.broadcast_in_dim %26568, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26570 = stablehlo.broadcast_in_dim %26567, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26571 = stablehlo.concatenate %26569, %26570, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26572 = stablehlo.abs %25459 : tensor<1x2048xf32> loc(#loc5307)
    %26573 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26574 = stablehlo.reduce(%26572 init: %26573) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26575 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26576 = stablehlo.maximum %26574, %26575 : tensor<f32> loc(#loc5309)
    %26577 = stablehlo.broadcast_in_dim %26576, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26578 = stablehlo.divide %25459, %26577 : tensor<1x2048xf32> loc(#loc5314)
    %26579 = stablehlo.multiply %26576, %26576 : tensor<f32> loc(#loc5310)
    %26580 = stablehlo.multiply %26578, %26578 : tensor<1x2048xf32> loc(#loc5310)
    %26581 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26582 = stablehlo.reduce(%26580 init: %26581) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26583 = stablehlo.multiply %26579, %26582 : tensor<f32> loc(#loc5310)
    %26584 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26585 = stablehlo.broadcast_in_dim %26584, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26586 = stablehlo.broadcast_in_dim %26583, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26587 = stablehlo.concatenate %26585, %26586, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26588 = stablehlo.abs %25460 : tensor<1x8192x2048xf32> loc(#loc5307)
    %26589 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26590 = stablehlo.reduce(%26588 init: %26589) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26591 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26592 = stablehlo.maximum %26590, %26591 : tensor<f32> loc(#loc5309)
    %26593 = stablehlo.broadcast_in_dim %26592, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %26594 = stablehlo.divide %25460, %26593 : tensor<1x8192x2048xf32> loc(#loc5314)
    %26595 = stablehlo.multiply %26592, %26592 : tensor<f32> loc(#loc5310)
    %26596 = stablehlo.multiply %26594, %26594 : tensor<1x8192x2048xf32> loc(#loc5310)
    %26597 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26598 = stablehlo.reduce(%26596 init: %26597) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26599 = stablehlo.multiply %26595, %26598 : tensor<f32> loc(#loc5310)
    %26600 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26601 = stablehlo.broadcast_in_dim %26600, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26602 = stablehlo.broadcast_in_dim %26599, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26603 = stablehlo.concatenate %26601, %26602, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26604 = stablehlo.abs %25461 : tensor<1x2048xf32> loc(#loc5307)
    %26605 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26606 = stablehlo.reduce(%26604 init: %26605) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26607 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26608 = stablehlo.maximum %26606, %26607 : tensor<f32> loc(#loc5309)
    %26609 = stablehlo.broadcast_in_dim %26608, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26610 = stablehlo.divide %25461, %26609 : tensor<1x2048xf32> loc(#loc5314)
    %26611 = stablehlo.multiply %26608, %26608 : tensor<f32> loc(#loc5310)
    %26612 = stablehlo.multiply %26610, %26610 : tensor<1x2048xf32> loc(#loc5310)
    %26613 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26614 = stablehlo.reduce(%26612 init: %26613) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26615 = stablehlo.multiply %26611, %26614 : tensor<f32> loc(#loc5310)
    %26616 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26617 = stablehlo.broadcast_in_dim %26616, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26618 = stablehlo.broadcast_in_dim %26615, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26619 = stablehlo.concatenate %26617, %26618, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26620 = stablehlo.abs %25462 : tensor<1x2048xf32> loc(#loc5307)
    %26621 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26622 = stablehlo.reduce(%26620 init: %26621) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26623 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26624 = stablehlo.maximum %26622, %26623 : tensor<f32> loc(#loc5309)
    %26625 = stablehlo.broadcast_in_dim %26624, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26626 = stablehlo.divide %25462, %26625 : tensor<1x2048xf32> loc(#loc5314)
    %26627 = stablehlo.multiply %26624, %26624 : tensor<f32> loc(#loc5310)
    %26628 = stablehlo.multiply %26626, %26626 : tensor<1x2048xf32> loc(#loc5310)
    %26629 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26630 = stablehlo.reduce(%26628 init: %26629) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26631 = stablehlo.multiply %26627, %26630 : tensor<f32> loc(#loc5310)
    %26632 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26633 = stablehlo.broadcast_in_dim %26632, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26634 = stablehlo.broadcast_in_dim %26631, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26635 = stablehlo.concatenate %26633, %26634, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26636 = stablehlo.abs %25463 : tensor<1x2048xf32> loc(#loc5307)
    %26637 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26638 = stablehlo.reduce(%26636 init: %26637) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26639 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26640 = stablehlo.maximum %26638, %26639 : tensor<f32> loc(#loc5309)
    %26641 = stablehlo.broadcast_in_dim %26640, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26642 = stablehlo.divide %25463, %26641 : tensor<1x2048xf32> loc(#loc5314)
    %26643 = stablehlo.multiply %26640, %26640 : tensor<f32> loc(#loc5310)
    %26644 = stablehlo.multiply %26642, %26642 : tensor<1x2048xf32> loc(#loc5310)
    %26645 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26646 = stablehlo.reduce(%26644 init: %26645) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26647 = stablehlo.multiply %26643, %26646 : tensor<f32> loc(#loc5310)
    %26648 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26649 = stablehlo.broadcast_in_dim %26648, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26650 = stablehlo.broadcast_in_dim %26647, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26651 = stablehlo.concatenate %26649, %26650, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26652 = stablehlo.abs %25464 : tensor<1x2048xf32> loc(#loc5307)
    %26653 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26654 = stablehlo.reduce(%26652 init: %26653) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26655 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26656 = stablehlo.maximum %26654, %26655 : tensor<f32> loc(#loc5309)
    %26657 = stablehlo.broadcast_in_dim %26656, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26658 = stablehlo.divide %25464, %26657 : tensor<1x2048xf32> loc(#loc5314)
    %26659 = stablehlo.multiply %26656, %26656 : tensor<f32> loc(#loc5310)
    %26660 = stablehlo.multiply %26658, %26658 : tensor<1x2048xf32> loc(#loc5310)
    %26661 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26662 = stablehlo.reduce(%26660 init: %26661) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26663 = stablehlo.multiply %26659, %26662 : tensor<f32> loc(#loc5310)
    %26664 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26665 = stablehlo.broadcast_in_dim %26664, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26666 = stablehlo.broadcast_in_dim %26663, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26667 = stablehlo.concatenate %26665, %26666, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26668 = stablehlo.abs %25465 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %26669 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %26670 = stablehlo.reduce(%26668 init: %26669) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %26671 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26672 = stablehlo.maximum %26670, %26671 : tensor<f32> loc(#loc5309)
    %26673 = stablehlo.broadcast_in_dim %26672, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26674 = stablehlo.divide %25465, %26673 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26675 = stablehlo.multiply %26672, %26672 : tensor<f32> loc(#loc5310)
    %26676 = stablehlo.multiply %26674, %26674 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %26677 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %26678 = stablehlo.reduce(%26676 init: %26677) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %26679 = stablehlo.multiply %26675, %26678 : tensor<f32> loc(#loc5310)
    %26680 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %26681 = stablehlo.broadcast_in_dim %26680, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26682 = stablehlo.broadcast_in_dim %26679, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26683 = stablehlo.concatenate %26681, %26682, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26684 = stablehlo.abs %25466 : tensor<1x64xf32> loc(#loc5307)
    %26685 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26686 = stablehlo.reduce(%26684 init: %26685) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26687 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26688 = stablehlo.maximum %26686, %26687 : tensor<f32> loc(#loc5309)
    %26689 = stablehlo.broadcast_in_dim %26688, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %26690 = stablehlo.divide %25466, %26689 : tensor<1x64xf32> loc(#loc5314)
    %26691 = stablehlo.multiply %26688, %26688 : tensor<f32> loc(#loc5310)
    %26692 = stablehlo.multiply %26690, %26690 : tensor<1x64xf32> loc(#loc5310)
    %26693 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26694 = stablehlo.reduce(%26692 init: %26693) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26695 = stablehlo.multiply %26691, %26694 : tensor<f32> loc(#loc5310)
    %26696 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %26697 = stablehlo.broadcast_in_dim %26696, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26698 = stablehlo.broadcast_in_dim %26695, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26699 = stablehlo.concatenate %26697, %26698, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26700 = stablehlo.abs %25467 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %26701 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %26702 = stablehlo.reduce(%26700 init: %26701) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %26703 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26704 = stablehlo.maximum %26702, %26703 : tensor<f32> loc(#loc5309)
    %26705 = stablehlo.broadcast_in_dim %26704, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %26706 = stablehlo.divide %25467, %26705 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %26707 = stablehlo.multiply %26704, %26704 : tensor<f32> loc(#loc5310)
    %26708 = stablehlo.multiply %26706, %26706 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %26709 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %26710 = stablehlo.reduce(%26708 init: %26709) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %26711 = stablehlo.multiply %26707, %26710 : tensor<f32> loc(#loc5310)
    %26712 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %26713 = stablehlo.broadcast_in_dim %26712, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26714 = stablehlo.broadcast_in_dim %26711, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26715 = stablehlo.concatenate %26713, %26714, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26716 = stablehlo.abs %25468 : tensor<1x8192xf32> loc(#loc5307)
    %26717 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26718 = stablehlo.reduce(%26716 init: %26717) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26719 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26720 = stablehlo.maximum %26718, %26719 : tensor<f32> loc(#loc5309)
    %26721 = stablehlo.broadcast_in_dim %26720, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %26722 = stablehlo.divide %25468, %26721 : tensor<1x8192xf32> loc(#loc5314)
    %26723 = stablehlo.multiply %26720, %26720 : tensor<f32> loc(#loc5310)
    %26724 = stablehlo.multiply %26722, %26722 : tensor<1x8192xf32> loc(#loc5310)
    %26725 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26726 = stablehlo.reduce(%26724 init: %26725) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26727 = stablehlo.multiply %26723, %26726 : tensor<f32> loc(#loc5310)
    %26728 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %26729 = stablehlo.broadcast_in_dim %26728, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26730 = stablehlo.broadcast_in_dim %26727, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26731 = stablehlo.concatenate %26729, %26730, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26732 = stablehlo.abs %25469 : tensor<1x2048x8192xf32> loc(#loc5307)
    %26733 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26734 = stablehlo.reduce(%26732 init: %26733) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26735 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26736 = stablehlo.maximum %26734, %26735 : tensor<f32> loc(#loc5309)
    %26737 = stablehlo.broadcast_in_dim %26736, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %26738 = stablehlo.divide %25469, %26737 : tensor<1x2048x8192xf32> loc(#loc5314)
    %26739 = stablehlo.multiply %26736, %26736 : tensor<f32> loc(#loc5310)
    %26740 = stablehlo.multiply %26738, %26738 : tensor<1x2048x8192xf32> loc(#loc5310)
    %26741 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26742 = stablehlo.reduce(%26740 init: %26741) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26743 = stablehlo.multiply %26739, %26742 : tensor<f32> loc(#loc5310)
    %26744 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26745 = stablehlo.broadcast_in_dim %26744, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26746 = stablehlo.broadcast_in_dim %26743, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26747 = stablehlo.concatenate %26745, %26746, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26748 = stablehlo.abs %25470 : tensor<1x2048xf32> loc(#loc5307)
    %26749 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26750 = stablehlo.reduce(%26748 init: %26749) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26751 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26752 = stablehlo.maximum %26750, %26751 : tensor<f32> loc(#loc5309)
    %26753 = stablehlo.broadcast_in_dim %26752, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26754 = stablehlo.divide %25470, %26753 : tensor<1x2048xf32> loc(#loc5314)
    %26755 = stablehlo.multiply %26752, %26752 : tensor<f32> loc(#loc5310)
    %26756 = stablehlo.multiply %26754, %26754 : tensor<1x2048xf32> loc(#loc5310)
    %26757 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26758 = stablehlo.reduce(%26756 init: %26757) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26759 = stablehlo.multiply %26755, %26758 : tensor<f32> loc(#loc5310)
    %26760 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26761 = stablehlo.broadcast_in_dim %26760, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26762 = stablehlo.broadcast_in_dim %26759, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26763 = stablehlo.concatenate %26761, %26762, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26764 = stablehlo.abs %25471 : tensor<1x8192x2048xf32> loc(#loc5307)
    %26765 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26766 = stablehlo.reduce(%26764 init: %26765) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26767 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26768 = stablehlo.maximum %26766, %26767 : tensor<f32> loc(#loc5309)
    %26769 = stablehlo.broadcast_in_dim %26768, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %26770 = stablehlo.divide %25471, %26769 : tensor<1x8192x2048xf32> loc(#loc5314)
    %26771 = stablehlo.multiply %26768, %26768 : tensor<f32> loc(#loc5310)
    %26772 = stablehlo.multiply %26770, %26770 : tensor<1x8192x2048xf32> loc(#loc5310)
    %26773 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26774 = stablehlo.reduce(%26772 init: %26773) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26775 = stablehlo.multiply %26771, %26774 : tensor<f32> loc(#loc5310)
    %26776 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26777 = stablehlo.broadcast_in_dim %26776, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26778 = stablehlo.broadcast_in_dim %26775, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26779 = stablehlo.concatenate %26777, %26778, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26780 = stablehlo.abs %25472 : tensor<1x2048xf32> loc(#loc5307)
    %26781 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26782 = stablehlo.reduce(%26780 init: %26781) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26783 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26784 = stablehlo.maximum %26782, %26783 : tensor<f32> loc(#loc5309)
    %26785 = stablehlo.broadcast_in_dim %26784, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26786 = stablehlo.divide %25472, %26785 : tensor<1x2048xf32> loc(#loc5314)
    %26787 = stablehlo.multiply %26784, %26784 : tensor<f32> loc(#loc5310)
    %26788 = stablehlo.multiply %26786, %26786 : tensor<1x2048xf32> loc(#loc5310)
    %26789 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26790 = stablehlo.reduce(%26788 init: %26789) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26791 = stablehlo.multiply %26787, %26790 : tensor<f32> loc(#loc5310)
    %26792 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26793 = stablehlo.broadcast_in_dim %26792, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26794 = stablehlo.broadcast_in_dim %26791, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26795 = stablehlo.concatenate %26793, %26794, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26796 = stablehlo.abs %25473 : tensor<1x2048xf32> loc(#loc5307)
    %26797 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26798 = stablehlo.reduce(%26796 init: %26797) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26799 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26800 = stablehlo.maximum %26798, %26799 : tensor<f32> loc(#loc5309)
    %26801 = stablehlo.broadcast_in_dim %26800, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26802 = stablehlo.divide %25473, %26801 : tensor<1x2048xf32> loc(#loc5314)
    %26803 = stablehlo.multiply %26800, %26800 : tensor<f32> loc(#loc5310)
    %26804 = stablehlo.multiply %26802, %26802 : tensor<1x2048xf32> loc(#loc5310)
    %26805 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26806 = stablehlo.reduce(%26804 init: %26805) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26807 = stablehlo.multiply %26803, %26806 : tensor<f32> loc(#loc5310)
    %26808 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26809 = stablehlo.broadcast_in_dim %26808, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26810 = stablehlo.broadcast_in_dim %26807, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26811 = stablehlo.concatenate %26809, %26810, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26812 = stablehlo.abs %25474 : tensor<1x2048xf32> loc(#loc5307)
    %26813 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26814 = stablehlo.reduce(%26812 init: %26813) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26815 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26816 = stablehlo.maximum %26814, %26815 : tensor<f32> loc(#loc5309)
    %26817 = stablehlo.broadcast_in_dim %26816, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26818 = stablehlo.divide %25474, %26817 : tensor<1x2048xf32> loc(#loc5314)
    %26819 = stablehlo.multiply %26816, %26816 : tensor<f32> loc(#loc5310)
    %26820 = stablehlo.multiply %26818, %26818 : tensor<1x2048xf32> loc(#loc5310)
    %26821 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26822 = stablehlo.reduce(%26820 init: %26821) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26823 = stablehlo.multiply %26819, %26822 : tensor<f32> loc(#loc5310)
    %26824 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26825 = stablehlo.broadcast_in_dim %26824, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26826 = stablehlo.broadcast_in_dim %26823, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26827 = stablehlo.concatenate %26825, %26826, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26828 = stablehlo.abs %25475 : tensor<1x2048xf32> loc(#loc5307)
    %26829 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26830 = stablehlo.reduce(%26828 init: %26829) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26831 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26832 = stablehlo.maximum %26830, %26831 : tensor<f32> loc(#loc5309)
    %26833 = stablehlo.broadcast_in_dim %26832, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26834 = stablehlo.divide %25475, %26833 : tensor<1x2048xf32> loc(#loc5314)
    %26835 = stablehlo.multiply %26832, %26832 : tensor<f32> loc(#loc5310)
    %26836 = stablehlo.multiply %26834, %26834 : tensor<1x2048xf32> loc(#loc5310)
    %26837 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26838 = stablehlo.reduce(%26836 init: %26837) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26839 = stablehlo.multiply %26835, %26838 : tensor<f32> loc(#loc5310)
    %26840 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26841 = stablehlo.broadcast_in_dim %26840, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26842 = stablehlo.broadcast_in_dim %26839, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26843 = stablehlo.concatenate %26841, %26842, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26844 = stablehlo.abs %25476 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %26845 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %26846 = stablehlo.reduce(%26844 init: %26845) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %26847 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26848 = stablehlo.maximum %26846, %26847 : tensor<f32> loc(#loc5309)
    %26849 = stablehlo.broadcast_in_dim %26848, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26850 = stablehlo.divide %25476, %26849 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %26851 = stablehlo.multiply %26848, %26848 : tensor<f32> loc(#loc5310)
    %26852 = stablehlo.multiply %26850, %26850 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %26853 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %26854 = stablehlo.reduce(%26852 init: %26853) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %26855 = stablehlo.multiply %26851, %26854 : tensor<f32> loc(#loc5310)
    %26856 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %26857 = stablehlo.broadcast_in_dim %26856, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26858 = stablehlo.broadcast_in_dim %26855, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26859 = stablehlo.concatenate %26857, %26858, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26860 = stablehlo.abs %25477 : tensor<1x64xf32> loc(#loc5307)
    %26861 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26862 = stablehlo.reduce(%26860 init: %26861) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26863 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26864 = stablehlo.maximum %26862, %26863 : tensor<f32> loc(#loc5309)
    %26865 = stablehlo.broadcast_in_dim %26864, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %26866 = stablehlo.divide %25477, %26865 : tensor<1x64xf32> loc(#loc5314)
    %26867 = stablehlo.multiply %26864, %26864 : tensor<f32> loc(#loc5310)
    %26868 = stablehlo.multiply %26866, %26866 : tensor<1x64xf32> loc(#loc5310)
    %26869 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26870 = stablehlo.reduce(%26868 init: %26869) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26871 = stablehlo.multiply %26867, %26870 : tensor<f32> loc(#loc5310)
    %26872 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %26873 = stablehlo.broadcast_in_dim %26872, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26874 = stablehlo.broadcast_in_dim %26871, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26875 = stablehlo.concatenate %26873, %26874, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26876 = stablehlo.abs %25478 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %26877 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %26878 = stablehlo.reduce(%26876 init: %26877) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %26879 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26880 = stablehlo.maximum %26878, %26879 : tensor<f32> loc(#loc5309)
    %26881 = stablehlo.broadcast_in_dim %26880, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %26882 = stablehlo.divide %25478, %26881 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %26883 = stablehlo.multiply %26880, %26880 : tensor<f32> loc(#loc5310)
    %26884 = stablehlo.multiply %26882, %26882 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %26885 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %26886 = stablehlo.reduce(%26884 init: %26885) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %26887 = stablehlo.multiply %26883, %26886 : tensor<f32> loc(#loc5310)
    %26888 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %26889 = stablehlo.broadcast_in_dim %26888, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26890 = stablehlo.broadcast_in_dim %26887, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26891 = stablehlo.concatenate %26889, %26890, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26892 = stablehlo.abs %25479 : tensor<1x8192xf32> loc(#loc5307)
    %26893 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26894 = stablehlo.reduce(%26892 init: %26893) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26895 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26896 = stablehlo.maximum %26894, %26895 : tensor<f32> loc(#loc5309)
    %26897 = stablehlo.broadcast_in_dim %26896, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %26898 = stablehlo.divide %25479, %26897 : tensor<1x8192xf32> loc(#loc5314)
    %26899 = stablehlo.multiply %26896, %26896 : tensor<f32> loc(#loc5310)
    %26900 = stablehlo.multiply %26898, %26898 : tensor<1x8192xf32> loc(#loc5310)
    %26901 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %26902 = stablehlo.reduce(%26900 init: %26901) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %26903 = stablehlo.multiply %26899, %26902 : tensor<f32> loc(#loc5310)
    %26904 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %26905 = stablehlo.broadcast_in_dim %26904, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26906 = stablehlo.broadcast_in_dim %26903, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26907 = stablehlo.concatenate %26905, %26906, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26908 = stablehlo.abs %25480 : tensor<1x2048x8192xf32> loc(#loc5307)
    %26909 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26910 = stablehlo.reduce(%26908 init: %26909) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26911 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26912 = stablehlo.maximum %26910, %26911 : tensor<f32> loc(#loc5309)
    %26913 = stablehlo.broadcast_in_dim %26912, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %26914 = stablehlo.divide %25480, %26913 : tensor<1x2048x8192xf32> loc(#loc5314)
    %26915 = stablehlo.multiply %26912, %26912 : tensor<f32> loc(#loc5310)
    %26916 = stablehlo.multiply %26914, %26914 : tensor<1x2048x8192xf32> loc(#loc5310)
    %26917 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26918 = stablehlo.reduce(%26916 init: %26917) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26919 = stablehlo.multiply %26915, %26918 : tensor<f32> loc(#loc5310)
    %26920 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26921 = stablehlo.broadcast_in_dim %26920, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26922 = stablehlo.broadcast_in_dim %26919, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26923 = stablehlo.concatenate %26921, %26922, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26924 = stablehlo.abs %25481 : tensor<1x2048xf32> loc(#loc5307)
    %26925 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26926 = stablehlo.reduce(%26924 init: %26925) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26927 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26928 = stablehlo.maximum %26926, %26927 : tensor<f32> loc(#loc5309)
    %26929 = stablehlo.broadcast_in_dim %26928, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26930 = stablehlo.divide %25481, %26929 : tensor<1x2048xf32> loc(#loc5314)
    %26931 = stablehlo.multiply %26928, %26928 : tensor<f32> loc(#loc5310)
    %26932 = stablehlo.multiply %26930, %26930 : tensor<1x2048xf32> loc(#loc5310)
    %26933 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26934 = stablehlo.reduce(%26932 init: %26933) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26935 = stablehlo.multiply %26931, %26934 : tensor<f32> loc(#loc5310)
    %26936 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26937 = stablehlo.broadcast_in_dim %26936, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26938 = stablehlo.broadcast_in_dim %26935, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26939 = stablehlo.concatenate %26937, %26938, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26940 = stablehlo.abs %25482 : tensor<1x8192x2048xf32> loc(#loc5307)
    %26941 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %26942 = stablehlo.reduce(%26940 init: %26941) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %26943 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26944 = stablehlo.maximum %26942, %26943 : tensor<f32> loc(#loc5309)
    %26945 = stablehlo.broadcast_in_dim %26944, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %26946 = stablehlo.divide %25482, %26945 : tensor<1x8192x2048xf32> loc(#loc5314)
    %26947 = stablehlo.multiply %26944, %26944 : tensor<f32> loc(#loc5310)
    %26948 = stablehlo.multiply %26946, %26946 : tensor<1x8192x2048xf32> loc(#loc5310)
    %26949 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %26950 = stablehlo.reduce(%26948 init: %26949) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %26951 = stablehlo.multiply %26947, %26950 : tensor<f32> loc(#loc5310)
    %26952 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %26953 = stablehlo.broadcast_in_dim %26952, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26954 = stablehlo.broadcast_in_dim %26951, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26955 = stablehlo.concatenate %26953, %26954, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26956 = stablehlo.abs %25483 : tensor<1x2048xf32> loc(#loc5307)
    %26957 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26958 = stablehlo.reduce(%26956 init: %26957) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26959 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26960 = stablehlo.maximum %26958, %26959 : tensor<f32> loc(#loc5309)
    %26961 = stablehlo.broadcast_in_dim %26960, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26962 = stablehlo.divide %25483, %26961 : tensor<1x2048xf32> loc(#loc5314)
    %26963 = stablehlo.multiply %26960, %26960 : tensor<f32> loc(#loc5310)
    %26964 = stablehlo.multiply %26962, %26962 : tensor<1x2048xf32> loc(#loc5310)
    %26965 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26966 = stablehlo.reduce(%26964 init: %26965) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26967 = stablehlo.multiply %26963, %26966 : tensor<f32> loc(#loc5310)
    %26968 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26969 = stablehlo.broadcast_in_dim %26968, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26970 = stablehlo.broadcast_in_dim %26967, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26971 = stablehlo.concatenate %26969, %26970, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26972 = stablehlo.abs %25484 : tensor<1x2048xf32> loc(#loc5307)
    %26973 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26974 = stablehlo.reduce(%26972 init: %26973) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26975 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26976 = stablehlo.maximum %26974, %26975 : tensor<f32> loc(#loc5309)
    %26977 = stablehlo.broadcast_in_dim %26976, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26978 = stablehlo.divide %25484, %26977 : tensor<1x2048xf32> loc(#loc5314)
    %26979 = stablehlo.multiply %26976, %26976 : tensor<f32> loc(#loc5310)
    %26980 = stablehlo.multiply %26978, %26978 : tensor<1x2048xf32> loc(#loc5310)
    %26981 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26982 = stablehlo.reduce(%26980 init: %26981) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26983 = stablehlo.multiply %26979, %26982 : tensor<f32> loc(#loc5310)
    %26984 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %26985 = stablehlo.broadcast_in_dim %26984, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26986 = stablehlo.broadcast_in_dim %26983, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %26987 = stablehlo.concatenate %26985, %26986, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %26988 = stablehlo.abs %25485 : tensor<1x2048xf32> loc(#loc5307)
    %26989 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %26990 = stablehlo.reduce(%26988 init: %26989) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %26991 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %26992 = stablehlo.maximum %26990, %26991 : tensor<f32> loc(#loc5309)
    %26993 = stablehlo.broadcast_in_dim %26992, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %26994 = stablehlo.divide %25485, %26993 : tensor<1x2048xf32> loc(#loc5314)
    %26995 = stablehlo.multiply %26992, %26992 : tensor<f32> loc(#loc5310)
    %26996 = stablehlo.multiply %26994, %26994 : tensor<1x2048xf32> loc(#loc5310)
    %26997 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %26998 = stablehlo.reduce(%26996 init: %26997) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %26999 = stablehlo.multiply %26995, %26998 : tensor<f32> loc(#loc5310)
    %27000 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27001 = stablehlo.broadcast_in_dim %27000, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27002 = stablehlo.broadcast_in_dim %26999, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27003 = stablehlo.concatenate %27001, %27002, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27004 = stablehlo.abs %25486 : tensor<1x2048xf32> loc(#loc5307)
    %27005 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27006 = stablehlo.reduce(%27004 init: %27005) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27007 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27008 = stablehlo.maximum %27006, %27007 : tensor<f32> loc(#loc5309)
    %27009 = stablehlo.broadcast_in_dim %27008, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27010 = stablehlo.divide %25486, %27009 : tensor<1x2048xf32> loc(#loc5314)
    %27011 = stablehlo.multiply %27008, %27008 : tensor<f32> loc(#loc5310)
    %27012 = stablehlo.multiply %27010, %27010 : tensor<1x2048xf32> loc(#loc5310)
    %27013 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27014 = stablehlo.reduce(%27012 init: %27013) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27015 = stablehlo.multiply %27011, %27014 : tensor<f32> loc(#loc5310)
    %27016 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27017 = stablehlo.broadcast_in_dim %27016, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27018 = stablehlo.broadcast_in_dim %27015, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27019 = stablehlo.concatenate %27017, %27018, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27020 = stablehlo.abs %25487 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %27021 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %27022 = stablehlo.reduce(%27020 init: %27021) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %27023 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27024 = stablehlo.maximum %27022, %27023 : tensor<f32> loc(#loc5309)
    %27025 = stablehlo.broadcast_in_dim %27024, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27026 = stablehlo.divide %25487, %27025 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27027 = stablehlo.multiply %27024, %27024 : tensor<f32> loc(#loc5310)
    %27028 = stablehlo.multiply %27026, %27026 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %27029 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %27030 = stablehlo.reduce(%27028 init: %27029) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %27031 = stablehlo.multiply %27027, %27030 : tensor<f32> loc(#loc5310)
    %27032 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %27033 = stablehlo.broadcast_in_dim %27032, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27034 = stablehlo.broadcast_in_dim %27031, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27035 = stablehlo.concatenate %27033, %27034, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27036 = stablehlo.abs %25488 : tensor<1x64xf32> loc(#loc5307)
    %27037 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27038 = stablehlo.reduce(%27036 init: %27037) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27039 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27040 = stablehlo.maximum %27038, %27039 : tensor<f32> loc(#loc5309)
    %27041 = stablehlo.broadcast_in_dim %27040, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %27042 = stablehlo.divide %25488, %27041 : tensor<1x64xf32> loc(#loc5314)
    %27043 = stablehlo.multiply %27040, %27040 : tensor<f32> loc(#loc5310)
    %27044 = stablehlo.multiply %27042, %27042 : tensor<1x64xf32> loc(#loc5310)
    %27045 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27046 = stablehlo.reduce(%27044 init: %27045) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27047 = stablehlo.multiply %27043, %27046 : tensor<f32> loc(#loc5310)
    %27048 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %27049 = stablehlo.broadcast_in_dim %27048, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27050 = stablehlo.broadcast_in_dim %27047, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27051 = stablehlo.concatenate %27049, %27050, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27052 = stablehlo.abs %25489 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %27053 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %27054 = stablehlo.reduce(%27052 init: %27053) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %27055 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27056 = stablehlo.maximum %27054, %27055 : tensor<f32> loc(#loc5309)
    %27057 = stablehlo.broadcast_in_dim %27056, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %27058 = stablehlo.divide %25489, %27057 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %27059 = stablehlo.multiply %27056, %27056 : tensor<f32> loc(#loc5310)
    %27060 = stablehlo.multiply %27058, %27058 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %27061 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %27062 = stablehlo.reduce(%27060 init: %27061) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %27063 = stablehlo.multiply %27059, %27062 : tensor<f32> loc(#loc5310)
    %27064 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %27065 = stablehlo.broadcast_in_dim %27064, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27066 = stablehlo.broadcast_in_dim %27063, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27067 = stablehlo.concatenate %27065, %27066, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27068 = stablehlo.abs %25490 : tensor<1x8192xf32> loc(#loc5307)
    %27069 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27070 = stablehlo.reduce(%27068 init: %27069) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27071 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27072 = stablehlo.maximum %27070, %27071 : tensor<f32> loc(#loc5309)
    %27073 = stablehlo.broadcast_in_dim %27072, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %27074 = stablehlo.divide %25490, %27073 : tensor<1x8192xf32> loc(#loc5314)
    %27075 = stablehlo.multiply %27072, %27072 : tensor<f32> loc(#loc5310)
    %27076 = stablehlo.multiply %27074, %27074 : tensor<1x8192xf32> loc(#loc5310)
    %27077 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27078 = stablehlo.reduce(%27076 init: %27077) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27079 = stablehlo.multiply %27075, %27078 : tensor<f32> loc(#loc5310)
    %27080 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %27081 = stablehlo.broadcast_in_dim %27080, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27082 = stablehlo.broadcast_in_dim %27079, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27083 = stablehlo.concatenate %27081, %27082, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27084 = stablehlo.abs %25491 : tensor<1x2048x8192xf32> loc(#loc5307)
    %27085 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27086 = stablehlo.reduce(%27084 init: %27085) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27087 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27088 = stablehlo.maximum %27086, %27087 : tensor<f32> loc(#loc5309)
    %27089 = stablehlo.broadcast_in_dim %27088, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %27090 = stablehlo.divide %25491, %27089 : tensor<1x2048x8192xf32> loc(#loc5314)
    %27091 = stablehlo.multiply %27088, %27088 : tensor<f32> loc(#loc5310)
    %27092 = stablehlo.multiply %27090, %27090 : tensor<1x2048x8192xf32> loc(#loc5310)
    %27093 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27094 = stablehlo.reduce(%27092 init: %27093) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27095 = stablehlo.multiply %27091, %27094 : tensor<f32> loc(#loc5310)
    %27096 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27097 = stablehlo.broadcast_in_dim %27096, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27098 = stablehlo.broadcast_in_dim %27095, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27099 = stablehlo.concatenate %27097, %27098, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27100 = stablehlo.abs %25492 : tensor<1x2048xf32> loc(#loc5307)
    %27101 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27102 = stablehlo.reduce(%27100 init: %27101) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27103 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27104 = stablehlo.maximum %27102, %27103 : tensor<f32> loc(#loc5309)
    %27105 = stablehlo.broadcast_in_dim %27104, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27106 = stablehlo.divide %25492, %27105 : tensor<1x2048xf32> loc(#loc5314)
    %27107 = stablehlo.multiply %27104, %27104 : tensor<f32> loc(#loc5310)
    %27108 = stablehlo.multiply %27106, %27106 : tensor<1x2048xf32> loc(#loc5310)
    %27109 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27110 = stablehlo.reduce(%27108 init: %27109) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27111 = stablehlo.multiply %27107, %27110 : tensor<f32> loc(#loc5310)
    %27112 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27113 = stablehlo.broadcast_in_dim %27112, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27114 = stablehlo.broadcast_in_dim %27111, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27115 = stablehlo.concatenate %27113, %27114, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27116 = stablehlo.abs %25493 : tensor<1x8192x2048xf32> loc(#loc5307)
    %27117 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27118 = stablehlo.reduce(%27116 init: %27117) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27119 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27120 = stablehlo.maximum %27118, %27119 : tensor<f32> loc(#loc5309)
    %27121 = stablehlo.broadcast_in_dim %27120, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %27122 = stablehlo.divide %25493, %27121 : tensor<1x8192x2048xf32> loc(#loc5314)
    %27123 = stablehlo.multiply %27120, %27120 : tensor<f32> loc(#loc5310)
    %27124 = stablehlo.multiply %27122, %27122 : tensor<1x8192x2048xf32> loc(#loc5310)
    %27125 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27126 = stablehlo.reduce(%27124 init: %27125) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27127 = stablehlo.multiply %27123, %27126 : tensor<f32> loc(#loc5310)
    %27128 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27129 = stablehlo.broadcast_in_dim %27128, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27130 = stablehlo.broadcast_in_dim %27127, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27131 = stablehlo.concatenate %27129, %27130, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27132 = stablehlo.abs %25494 : tensor<1x2048xf32> loc(#loc5307)
    %27133 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27134 = stablehlo.reduce(%27132 init: %27133) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27135 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27136 = stablehlo.maximum %27134, %27135 : tensor<f32> loc(#loc5309)
    %27137 = stablehlo.broadcast_in_dim %27136, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27138 = stablehlo.divide %25494, %27137 : tensor<1x2048xf32> loc(#loc5314)
    %27139 = stablehlo.multiply %27136, %27136 : tensor<f32> loc(#loc5310)
    %27140 = stablehlo.multiply %27138, %27138 : tensor<1x2048xf32> loc(#loc5310)
    %27141 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27142 = stablehlo.reduce(%27140 init: %27141) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27143 = stablehlo.multiply %27139, %27142 : tensor<f32> loc(#loc5310)
    %27144 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27145 = stablehlo.broadcast_in_dim %27144, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27146 = stablehlo.broadcast_in_dim %27143, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27147 = stablehlo.concatenate %27145, %27146, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27148 = stablehlo.abs %25495 : tensor<1x2048xf32> loc(#loc5307)
    %27149 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27150 = stablehlo.reduce(%27148 init: %27149) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27151 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27152 = stablehlo.maximum %27150, %27151 : tensor<f32> loc(#loc5309)
    %27153 = stablehlo.broadcast_in_dim %27152, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27154 = stablehlo.divide %25495, %27153 : tensor<1x2048xf32> loc(#loc5314)
    %27155 = stablehlo.multiply %27152, %27152 : tensor<f32> loc(#loc5310)
    %27156 = stablehlo.multiply %27154, %27154 : tensor<1x2048xf32> loc(#loc5310)
    %27157 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27158 = stablehlo.reduce(%27156 init: %27157) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27159 = stablehlo.multiply %27155, %27158 : tensor<f32> loc(#loc5310)
    %27160 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27161 = stablehlo.broadcast_in_dim %27160, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27162 = stablehlo.broadcast_in_dim %27159, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27163 = stablehlo.concatenate %27161, %27162, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27164 = stablehlo.abs %25496 : tensor<1x2048xf32> loc(#loc5307)
    %27165 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27166 = stablehlo.reduce(%27164 init: %27165) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27167 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27168 = stablehlo.maximum %27166, %27167 : tensor<f32> loc(#loc5309)
    %27169 = stablehlo.broadcast_in_dim %27168, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27170 = stablehlo.divide %25496, %27169 : tensor<1x2048xf32> loc(#loc5314)
    %27171 = stablehlo.multiply %27168, %27168 : tensor<f32> loc(#loc5310)
    %27172 = stablehlo.multiply %27170, %27170 : tensor<1x2048xf32> loc(#loc5310)
    %27173 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27174 = stablehlo.reduce(%27172 init: %27173) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27175 = stablehlo.multiply %27171, %27174 : tensor<f32> loc(#loc5310)
    %27176 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27177 = stablehlo.broadcast_in_dim %27176, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27178 = stablehlo.broadcast_in_dim %27175, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27179 = stablehlo.concatenate %27177, %27178, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27180 = stablehlo.abs %25497 : tensor<1x2048xf32> loc(#loc5307)
    %27181 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27182 = stablehlo.reduce(%27180 init: %27181) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27183 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27184 = stablehlo.maximum %27182, %27183 : tensor<f32> loc(#loc5309)
    %27185 = stablehlo.broadcast_in_dim %27184, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27186 = stablehlo.divide %25497, %27185 : tensor<1x2048xf32> loc(#loc5314)
    %27187 = stablehlo.multiply %27184, %27184 : tensor<f32> loc(#loc5310)
    %27188 = stablehlo.multiply %27186, %27186 : tensor<1x2048xf32> loc(#loc5310)
    %27189 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27190 = stablehlo.reduce(%27188 init: %27189) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27191 = stablehlo.multiply %27187, %27190 : tensor<f32> loc(#loc5310)
    %27192 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27193 = stablehlo.broadcast_in_dim %27192, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27194 = stablehlo.broadcast_in_dim %27191, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27195 = stablehlo.concatenate %27193, %27194, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27196 = stablehlo.abs %25498 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %27197 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %27198 = stablehlo.reduce(%27196 init: %27197) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %27199 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27200 = stablehlo.maximum %27198, %27199 : tensor<f32> loc(#loc5309)
    %27201 = stablehlo.broadcast_in_dim %27200, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27202 = stablehlo.divide %25498, %27201 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27203 = stablehlo.multiply %27200, %27200 : tensor<f32> loc(#loc5310)
    %27204 = stablehlo.multiply %27202, %27202 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %27205 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %27206 = stablehlo.reduce(%27204 init: %27205) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %27207 = stablehlo.multiply %27203, %27206 : tensor<f32> loc(#loc5310)
    %27208 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %27209 = stablehlo.broadcast_in_dim %27208, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27210 = stablehlo.broadcast_in_dim %27207, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27211 = stablehlo.concatenate %27209, %27210, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27212 = stablehlo.abs %25499 : tensor<1x64xf32> loc(#loc5307)
    %27213 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27214 = stablehlo.reduce(%27212 init: %27213) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27215 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27216 = stablehlo.maximum %27214, %27215 : tensor<f32> loc(#loc5309)
    %27217 = stablehlo.broadcast_in_dim %27216, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %27218 = stablehlo.divide %25499, %27217 : tensor<1x64xf32> loc(#loc5314)
    %27219 = stablehlo.multiply %27216, %27216 : tensor<f32> loc(#loc5310)
    %27220 = stablehlo.multiply %27218, %27218 : tensor<1x64xf32> loc(#loc5310)
    %27221 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27222 = stablehlo.reduce(%27220 init: %27221) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27223 = stablehlo.multiply %27219, %27222 : tensor<f32> loc(#loc5310)
    %27224 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %27225 = stablehlo.broadcast_in_dim %27224, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27226 = stablehlo.broadcast_in_dim %27223, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27227 = stablehlo.concatenate %27225, %27226, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27228 = stablehlo.abs %25500 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %27229 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %27230 = stablehlo.reduce(%27228 init: %27229) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %27231 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27232 = stablehlo.maximum %27230, %27231 : tensor<f32> loc(#loc5309)
    %27233 = stablehlo.broadcast_in_dim %27232, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %27234 = stablehlo.divide %25500, %27233 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %27235 = stablehlo.multiply %27232, %27232 : tensor<f32> loc(#loc5310)
    %27236 = stablehlo.multiply %27234, %27234 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %27237 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %27238 = stablehlo.reduce(%27236 init: %27237) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %27239 = stablehlo.multiply %27235, %27238 : tensor<f32> loc(#loc5310)
    %27240 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %27241 = stablehlo.broadcast_in_dim %27240, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27242 = stablehlo.broadcast_in_dim %27239, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27243 = stablehlo.concatenate %27241, %27242, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27244 = stablehlo.abs %25501 : tensor<1x8192xf32> loc(#loc5307)
    %27245 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27246 = stablehlo.reduce(%27244 init: %27245) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27247 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27248 = stablehlo.maximum %27246, %27247 : tensor<f32> loc(#loc5309)
    %27249 = stablehlo.broadcast_in_dim %27248, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %27250 = stablehlo.divide %25501, %27249 : tensor<1x8192xf32> loc(#loc5314)
    %27251 = stablehlo.multiply %27248, %27248 : tensor<f32> loc(#loc5310)
    %27252 = stablehlo.multiply %27250, %27250 : tensor<1x8192xf32> loc(#loc5310)
    %27253 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27254 = stablehlo.reduce(%27252 init: %27253) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27255 = stablehlo.multiply %27251, %27254 : tensor<f32> loc(#loc5310)
    %27256 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %27257 = stablehlo.broadcast_in_dim %27256, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27258 = stablehlo.broadcast_in_dim %27255, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27259 = stablehlo.concatenate %27257, %27258, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27260 = stablehlo.abs %25502 : tensor<1x2048x8192xf32> loc(#loc5307)
    %27261 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27262 = stablehlo.reduce(%27260 init: %27261) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27263 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27264 = stablehlo.maximum %27262, %27263 : tensor<f32> loc(#loc5309)
    %27265 = stablehlo.broadcast_in_dim %27264, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %27266 = stablehlo.divide %25502, %27265 : tensor<1x2048x8192xf32> loc(#loc5314)
    %27267 = stablehlo.multiply %27264, %27264 : tensor<f32> loc(#loc5310)
    %27268 = stablehlo.multiply %27266, %27266 : tensor<1x2048x8192xf32> loc(#loc5310)
    %27269 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27270 = stablehlo.reduce(%27268 init: %27269) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27271 = stablehlo.multiply %27267, %27270 : tensor<f32> loc(#loc5310)
    %27272 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27273 = stablehlo.broadcast_in_dim %27272, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27274 = stablehlo.broadcast_in_dim %27271, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27275 = stablehlo.concatenate %27273, %27274, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27276 = stablehlo.abs %25503 : tensor<1x2048xf32> loc(#loc5307)
    %27277 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27278 = stablehlo.reduce(%27276 init: %27277) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27279 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27280 = stablehlo.maximum %27278, %27279 : tensor<f32> loc(#loc5309)
    %27281 = stablehlo.broadcast_in_dim %27280, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27282 = stablehlo.divide %25503, %27281 : tensor<1x2048xf32> loc(#loc5314)
    %27283 = stablehlo.multiply %27280, %27280 : tensor<f32> loc(#loc5310)
    %27284 = stablehlo.multiply %27282, %27282 : tensor<1x2048xf32> loc(#loc5310)
    %27285 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27286 = stablehlo.reduce(%27284 init: %27285) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27287 = stablehlo.multiply %27283, %27286 : tensor<f32> loc(#loc5310)
    %27288 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27289 = stablehlo.broadcast_in_dim %27288, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27290 = stablehlo.broadcast_in_dim %27287, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27291 = stablehlo.concatenate %27289, %27290, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27292 = stablehlo.abs %25504 : tensor<1x8192x2048xf32> loc(#loc5307)
    %27293 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27294 = stablehlo.reduce(%27292 init: %27293) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27295 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27296 = stablehlo.maximum %27294, %27295 : tensor<f32> loc(#loc5309)
    %27297 = stablehlo.broadcast_in_dim %27296, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %27298 = stablehlo.divide %25504, %27297 : tensor<1x8192x2048xf32> loc(#loc5314)
    %27299 = stablehlo.multiply %27296, %27296 : tensor<f32> loc(#loc5310)
    %27300 = stablehlo.multiply %27298, %27298 : tensor<1x8192x2048xf32> loc(#loc5310)
    %27301 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27302 = stablehlo.reduce(%27300 init: %27301) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27303 = stablehlo.multiply %27299, %27302 : tensor<f32> loc(#loc5310)
    %27304 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27305 = stablehlo.broadcast_in_dim %27304, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27306 = stablehlo.broadcast_in_dim %27303, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27307 = stablehlo.concatenate %27305, %27306, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27308 = stablehlo.abs %25505 : tensor<1x2048xf32> loc(#loc5307)
    %27309 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27310 = stablehlo.reduce(%27308 init: %27309) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27311 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27312 = stablehlo.maximum %27310, %27311 : tensor<f32> loc(#loc5309)
    %27313 = stablehlo.broadcast_in_dim %27312, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27314 = stablehlo.divide %25505, %27313 : tensor<1x2048xf32> loc(#loc5314)
    %27315 = stablehlo.multiply %27312, %27312 : tensor<f32> loc(#loc5310)
    %27316 = stablehlo.multiply %27314, %27314 : tensor<1x2048xf32> loc(#loc5310)
    %27317 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27318 = stablehlo.reduce(%27316 init: %27317) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27319 = stablehlo.multiply %27315, %27318 : tensor<f32> loc(#loc5310)
    %27320 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27321 = stablehlo.broadcast_in_dim %27320, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27322 = stablehlo.broadcast_in_dim %27319, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27323 = stablehlo.concatenate %27321, %27322, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27324 = stablehlo.abs %25506 : tensor<1x2048xf32> loc(#loc5307)
    %27325 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27326 = stablehlo.reduce(%27324 init: %27325) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27327 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27328 = stablehlo.maximum %27326, %27327 : tensor<f32> loc(#loc5309)
    %27329 = stablehlo.broadcast_in_dim %27328, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27330 = stablehlo.divide %25506, %27329 : tensor<1x2048xf32> loc(#loc5314)
    %27331 = stablehlo.multiply %27328, %27328 : tensor<f32> loc(#loc5310)
    %27332 = stablehlo.multiply %27330, %27330 : tensor<1x2048xf32> loc(#loc5310)
    %27333 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27334 = stablehlo.reduce(%27332 init: %27333) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27335 = stablehlo.multiply %27331, %27334 : tensor<f32> loc(#loc5310)
    %27336 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27337 = stablehlo.broadcast_in_dim %27336, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27338 = stablehlo.broadcast_in_dim %27335, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27339 = stablehlo.concatenate %27337, %27338, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27340 = stablehlo.abs %25507 : tensor<1x2048xf32> loc(#loc5307)
    %27341 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27342 = stablehlo.reduce(%27340 init: %27341) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27343 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27344 = stablehlo.maximum %27342, %27343 : tensor<f32> loc(#loc5309)
    %27345 = stablehlo.broadcast_in_dim %27344, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27346 = stablehlo.divide %25507, %27345 : tensor<1x2048xf32> loc(#loc5314)
    %27347 = stablehlo.multiply %27344, %27344 : tensor<f32> loc(#loc5310)
    %27348 = stablehlo.multiply %27346, %27346 : tensor<1x2048xf32> loc(#loc5310)
    %27349 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27350 = stablehlo.reduce(%27348 init: %27349) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27351 = stablehlo.multiply %27347, %27350 : tensor<f32> loc(#loc5310)
    %27352 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27353 = stablehlo.broadcast_in_dim %27352, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27354 = stablehlo.broadcast_in_dim %27351, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27355 = stablehlo.concatenate %27353, %27354, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27356 = stablehlo.abs %25508 : tensor<1x2048xf32> loc(#loc5307)
    %27357 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27358 = stablehlo.reduce(%27356 init: %27357) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27359 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27360 = stablehlo.maximum %27358, %27359 : tensor<f32> loc(#loc5309)
    %27361 = stablehlo.broadcast_in_dim %27360, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27362 = stablehlo.divide %25508, %27361 : tensor<1x2048xf32> loc(#loc5314)
    %27363 = stablehlo.multiply %27360, %27360 : tensor<f32> loc(#loc5310)
    %27364 = stablehlo.multiply %27362, %27362 : tensor<1x2048xf32> loc(#loc5310)
    %27365 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27366 = stablehlo.reduce(%27364 init: %27365) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27367 = stablehlo.multiply %27363, %27366 : tensor<f32> loc(#loc5310)
    %27368 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27369 = stablehlo.broadcast_in_dim %27368, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27370 = stablehlo.broadcast_in_dim %27367, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27371 = stablehlo.concatenate %27369, %27370, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27372 = stablehlo.abs %25509 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %27373 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %27374 = stablehlo.reduce(%27372 init: %27373) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %27375 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27376 = stablehlo.maximum %27374, %27375 : tensor<f32> loc(#loc5309)
    %27377 = stablehlo.broadcast_in_dim %27376, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27378 = stablehlo.divide %25509, %27377 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27379 = stablehlo.multiply %27376, %27376 : tensor<f32> loc(#loc5310)
    %27380 = stablehlo.multiply %27378, %27378 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %27381 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %27382 = stablehlo.reduce(%27380 init: %27381) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %27383 = stablehlo.multiply %27379, %27382 : tensor<f32> loc(#loc5310)
    %27384 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %27385 = stablehlo.broadcast_in_dim %27384, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27386 = stablehlo.broadcast_in_dim %27383, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27387 = stablehlo.concatenate %27385, %27386, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27388 = stablehlo.abs %25510 : tensor<1x64xf32> loc(#loc5307)
    %27389 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27390 = stablehlo.reduce(%27388 init: %27389) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27391 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27392 = stablehlo.maximum %27390, %27391 : tensor<f32> loc(#loc5309)
    %27393 = stablehlo.broadcast_in_dim %27392, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %27394 = stablehlo.divide %25510, %27393 : tensor<1x64xf32> loc(#loc5314)
    %27395 = stablehlo.multiply %27392, %27392 : tensor<f32> loc(#loc5310)
    %27396 = stablehlo.multiply %27394, %27394 : tensor<1x64xf32> loc(#loc5310)
    %27397 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27398 = stablehlo.reduce(%27396 init: %27397) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27399 = stablehlo.multiply %27395, %27398 : tensor<f32> loc(#loc5310)
    %27400 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %27401 = stablehlo.broadcast_in_dim %27400, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27402 = stablehlo.broadcast_in_dim %27399, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27403 = stablehlo.concatenate %27401, %27402, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27404 = stablehlo.abs %25511 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %27405 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %27406 = stablehlo.reduce(%27404 init: %27405) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %27407 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27408 = stablehlo.maximum %27406, %27407 : tensor<f32> loc(#loc5309)
    %27409 = stablehlo.broadcast_in_dim %27408, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %27410 = stablehlo.divide %25511, %27409 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %27411 = stablehlo.multiply %27408, %27408 : tensor<f32> loc(#loc5310)
    %27412 = stablehlo.multiply %27410, %27410 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %27413 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %27414 = stablehlo.reduce(%27412 init: %27413) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %27415 = stablehlo.multiply %27411, %27414 : tensor<f32> loc(#loc5310)
    %27416 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %27417 = stablehlo.broadcast_in_dim %27416, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27418 = stablehlo.broadcast_in_dim %27415, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27419 = stablehlo.concatenate %27417, %27418, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27420 = stablehlo.abs %25512 : tensor<1x8192xf32> loc(#loc5307)
    %27421 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27422 = stablehlo.reduce(%27420 init: %27421) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27423 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27424 = stablehlo.maximum %27422, %27423 : tensor<f32> loc(#loc5309)
    %27425 = stablehlo.broadcast_in_dim %27424, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %27426 = stablehlo.divide %25512, %27425 : tensor<1x8192xf32> loc(#loc5314)
    %27427 = stablehlo.multiply %27424, %27424 : tensor<f32> loc(#loc5310)
    %27428 = stablehlo.multiply %27426, %27426 : tensor<1x8192xf32> loc(#loc5310)
    %27429 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27430 = stablehlo.reduce(%27428 init: %27429) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27431 = stablehlo.multiply %27427, %27430 : tensor<f32> loc(#loc5310)
    %27432 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %27433 = stablehlo.broadcast_in_dim %27432, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27434 = stablehlo.broadcast_in_dim %27431, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27435 = stablehlo.concatenate %27433, %27434, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27436 = stablehlo.abs %25513 : tensor<1x2048x8192xf32> loc(#loc5307)
    %27437 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27438 = stablehlo.reduce(%27436 init: %27437) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27439 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27440 = stablehlo.maximum %27438, %27439 : tensor<f32> loc(#loc5309)
    %27441 = stablehlo.broadcast_in_dim %27440, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %27442 = stablehlo.divide %25513, %27441 : tensor<1x2048x8192xf32> loc(#loc5314)
    %27443 = stablehlo.multiply %27440, %27440 : tensor<f32> loc(#loc5310)
    %27444 = stablehlo.multiply %27442, %27442 : tensor<1x2048x8192xf32> loc(#loc5310)
    %27445 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27446 = stablehlo.reduce(%27444 init: %27445) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27447 = stablehlo.multiply %27443, %27446 : tensor<f32> loc(#loc5310)
    %27448 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27449 = stablehlo.broadcast_in_dim %27448, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27450 = stablehlo.broadcast_in_dim %27447, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27451 = stablehlo.concatenate %27449, %27450, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27452 = stablehlo.abs %25514 : tensor<1x2048xf32> loc(#loc5307)
    %27453 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27454 = stablehlo.reduce(%27452 init: %27453) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27455 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27456 = stablehlo.maximum %27454, %27455 : tensor<f32> loc(#loc5309)
    %27457 = stablehlo.broadcast_in_dim %27456, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27458 = stablehlo.divide %25514, %27457 : tensor<1x2048xf32> loc(#loc5314)
    %27459 = stablehlo.multiply %27456, %27456 : tensor<f32> loc(#loc5310)
    %27460 = stablehlo.multiply %27458, %27458 : tensor<1x2048xf32> loc(#loc5310)
    %27461 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27462 = stablehlo.reduce(%27460 init: %27461) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27463 = stablehlo.multiply %27459, %27462 : tensor<f32> loc(#loc5310)
    %27464 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27465 = stablehlo.broadcast_in_dim %27464, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27466 = stablehlo.broadcast_in_dim %27463, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27467 = stablehlo.concatenate %27465, %27466, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27468 = stablehlo.abs %25515 : tensor<1x8192x2048xf32> loc(#loc5307)
    %27469 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27470 = stablehlo.reduce(%27468 init: %27469) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27471 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27472 = stablehlo.maximum %27470, %27471 : tensor<f32> loc(#loc5309)
    %27473 = stablehlo.broadcast_in_dim %27472, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %27474 = stablehlo.divide %25515, %27473 : tensor<1x8192x2048xf32> loc(#loc5314)
    %27475 = stablehlo.multiply %27472, %27472 : tensor<f32> loc(#loc5310)
    %27476 = stablehlo.multiply %27474, %27474 : tensor<1x8192x2048xf32> loc(#loc5310)
    %27477 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27478 = stablehlo.reduce(%27476 init: %27477) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27479 = stablehlo.multiply %27475, %27478 : tensor<f32> loc(#loc5310)
    %27480 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27481 = stablehlo.broadcast_in_dim %27480, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27482 = stablehlo.broadcast_in_dim %27479, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27483 = stablehlo.concatenate %27481, %27482, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27484 = stablehlo.abs %25516 : tensor<1x2048xf32> loc(#loc5307)
    %27485 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27486 = stablehlo.reduce(%27484 init: %27485) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27487 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27488 = stablehlo.maximum %27486, %27487 : tensor<f32> loc(#loc5309)
    %27489 = stablehlo.broadcast_in_dim %27488, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27490 = stablehlo.divide %25516, %27489 : tensor<1x2048xf32> loc(#loc5314)
    %27491 = stablehlo.multiply %27488, %27488 : tensor<f32> loc(#loc5310)
    %27492 = stablehlo.multiply %27490, %27490 : tensor<1x2048xf32> loc(#loc5310)
    %27493 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27494 = stablehlo.reduce(%27492 init: %27493) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27495 = stablehlo.multiply %27491, %27494 : tensor<f32> loc(#loc5310)
    %27496 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27497 = stablehlo.broadcast_in_dim %27496, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27498 = stablehlo.broadcast_in_dim %27495, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27499 = stablehlo.concatenate %27497, %27498, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27500 = stablehlo.abs %25517 : tensor<1x2048xf32> loc(#loc5307)
    %27501 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27502 = stablehlo.reduce(%27500 init: %27501) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27503 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27504 = stablehlo.maximum %27502, %27503 : tensor<f32> loc(#loc5309)
    %27505 = stablehlo.broadcast_in_dim %27504, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27506 = stablehlo.divide %25517, %27505 : tensor<1x2048xf32> loc(#loc5314)
    %27507 = stablehlo.multiply %27504, %27504 : tensor<f32> loc(#loc5310)
    %27508 = stablehlo.multiply %27506, %27506 : tensor<1x2048xf32> loc(#loc5310)
    %27509 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27510 = stablehlo.reduce(%27508 init: %27509) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27511 = stablehlo.multiply %27507, %27510 : tensor<f32> loc(#loc5310)
    %27512 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27513 = stablehlo.broadcast_in_dim %27512, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27514 = stablehlo.broadcast_in_dim %27511, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27515 = stablehlo.concatenate %27513, %27514, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27516 = stablehlo.abs %25518 : tensor<1x2048xf32> loc(#loc5307)
    %27517 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27518 = stablehlo.reduce(%27516 init: %27517) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27519 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27520 = stablehlo.maximum %27518, %27519 : tensor<f32> loc(#loc5309)
    %27521 = stablehlo.broadcast_in_dim %27520, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27522 = stablehlo.divide %25518, %27521 : tensor<1x2048xf32> loc(#loc5314)
    %27523 = stablehlo.multiply %27520, %27520 : tensor<f32> loc(#loc5310)
    %27524 = stablehlo.multiply %27522, %27522 : tensor<1x2048xf32> loc(#loc5310)
    %27525 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27526 = stablehlo.reduce(%27524 init: %27525) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27527 = stablehlo.multiply %27523, %27526 : tensor<f32> loc(#loc5310)
    %27528 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27529 = stablehlo.broadcast_in_dim %27528, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27530 = stablehlo.broadcast_in_dim %27527, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27531 = stablehlo.concatenate %27529, %27530, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27532 = stablehlo.abs %25519 : tensor<1x2048xf32> loc(#loc5307)
    %27533 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27534 = stablehlo.reduce(%27532 init: %27533) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27535 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27536 = stablehlo.maximum %27534, %27535 : tensor<f32> loc(#loc5309)
    %27537 = stablehlo.broadcast_in_dim %27536, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27538 = stablehlo.divide %25519, %27537 : tensor<1x2048xf32> loc(#loc5314)
    %27539 = stablehlo.multiply %27536, %27536 : tensor<f32> loc(#loc5310)
    %27540 = stablehlo.multiply %27538, %27538 : tensor<1x2048xf32> loc(#loc5310)
    %27541 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27542 = stablehlo.reduce(%27540 init: %27541) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27543 = stablehlo.multiply %27539, %27542 : tensor<f32> loc(#loc5310)
    %27544 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27545 = stablehlo.broadcast_in_dim %27544, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27546 = stablehlo.broadcast_in_dim %27543, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27547 = stablehlo.concatenate %27545, %27546, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27548 = stablehlo.abs %25520 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %27549 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %27550 = stablehlo.reduce(%27548 init: %27549) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %27551 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27552 = stablehlo.maximum %27550, %27551 : tensor<f32> loc(#loc5309)
    %27553 = stablehlo.broadcast_in_dim %27552, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27554 = stablehlo.divide %25520, %27553 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27555 = stablehlo.multiply %27552, %27552 : tensor<f32> loc(#loc5310)
    %27556 = stablehlo.multiply %27554, %27554 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %27557 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %27558 = stablehlo.reduce(%27556 init: %27557) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %27559 = stablehlo.multiply %27555, %27558 : tensor<f32> loc(#loc5310)
    %27560 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %27561 = stablehlo.broadcast_in_dim %27560, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27562 = stablehlo.broadcast_in_dim %27559, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27563 = stablehlo.concatenate %27561, %27562, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27564 = stablehlo.abs %25521 : tensor<1x64xf32> loc(#loc5307)
    %27565 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27566 = stablehlo.reduce(%27564 init: %27565) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27567 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27568 = stablehlo.maximum %27566, %27567 : tensor<f32> loc(#loc5309)
    %27569 = stablehlo.broadcast_in_dim %27568, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %27570 = stablehlo.divide %25521, %27569 : tensor<1x64xf32> loc(#loc5314)
    %27571 = stablehlo.multiply %27568, %27568 : tensor<f32> loc(#loc5310)
    %27572 = stablehlo.multiply %27570, %27570 : tensor<1x64xf32> loc(#loc5310)
    %27573 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27574 = stablehlo.reduce(%27572 init: %27573) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27575 = stablehlo.multiply %27571, %27574 : tensor<f32> loc(#loc5310)
    %27576 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %27577 = stablehlo.broadcast_in_dim %27576, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27578 = stablehlo.broadcast_in_dim %27575, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27579 = stablehlo.concatenate %27577, %27578, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27580 = stablehlo.abs %25522 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %27581 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %27582 = stablehlo.reduce(%27580 init: %27581) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %27583 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27584 = stablehlo.maximum %27582, %27583 : tensor<f32> loc(#loc5309)
    %27585 = stablehlo.broadcast_in_dim %27584, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %27586 = stablehlo.divide %25522, %27585 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %27587 = stablehlo.multiply %27584, %27584 : tensor<f32> loc(#loc5310)
    %27588 = stablehlo.multiply %27586, %27586 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %27589 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %27590 = stablehlo.reduce(%27588 init: %27589) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %27591 = stablehlo.multiply %27587, %27590 : tensor<f32> loc(#loc5310)
    %27592 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %27593 = stablehlo.broadcast_in_dim %27592, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27594 = stablehlo.broadcast_in_dim %27591, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27595 = stablehlo.concatenate %27593, %27594, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27596 = stablehlo.abs %25523 : tensor<1x8192xf32> loc(#loc5307)
    %27597 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27598 = stablehlo.reduce(%27596 init: %27597) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27599 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27600 = stablehlo.maximum %27598, %27599 : tensor<f32> loc(#loc5309)
    %27601 = stablehlo.broadcast_in_dim %27600, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %27602 = stablehlo.divide %25523, %27601 : tensor<1x8192xf32> loc(#loc5314)
    %27603 = stablehlo.multiply %27600, %27600 : tensor<f32> loc(#loc5310)
    %27604 = stablehlo.multiply %27602, %27602 : tensor<1x8192xf32> loc(#loc5310)
    %27605 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27606 = stablehlo.reduce(%27604 init: %27605) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27607 = stablehlo.multiply %27603, %27606 : tensor<f32> loc(#loc5310)
    %27608 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %27609 = stablehlo.broadcast_in_dim %27608, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27610 = stablehlo.broadcast_in_dim %27607, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27611 = stablehlo.concatenate %27609, %27610, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27612 = stablehlo.abs %25524 : tensor<1x2048x8192xf32> loc(#loc5307)
    %27613 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27614 = stablehlo.reduce(%27612 init: %27613) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27615 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27616 = stablehlo.maximum %27614, %27615 : tensor<f32> loc(#loc5309)
    %27617 = stablehlo.broadcast_in_dim %27616, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %27618 = stablehlo.divide %25524, %27617 : tensor<1x2048x8192xf32> loc(#loc5314)
    %27619 = stablehlo.multiply %27616, %27616 : tensor<f32> loc(#loc5310)
    %27620 = stablehlo.multiply %27618, %27618 : tensor<1x2048x8192xf32> loc(#loc5310)
    %27621 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27622 = stablehlo.reduce(%27620 init: %27621) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27623 = stablehlo.multiply %27619, %27622 : tensor<f32> loc(#loc5310)
    %27624 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27625 = stablehlo.broadcast_in_dim %27624, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27626 = stablehlo.broadcast_in_dim %27623, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27627 = stablehlo.concatenate %27625, %27626, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27628 = stablehlo.abs %25525 : tensor<1x2048xf32> loc(#loc5307)
    %27629 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27630 = stablehlo.reduce(%27628 init: %27629) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27631 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27632 = stablehlo.maximum %27630, %27631 : tensor<f32> loc(#loc5309)
    %27633 = stablehlo.broadcast_in_dim %27632, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27634 = stablehlo.divide %25525, %27633 : tensor<1x2048xf32> loc(#loc5314)
    %27635 = stablehlo.multiply %27632, %27632 : tensor<f32> loc(#loc5310)
    %27636 = stablehlo.multiply %27634, %27634 : tensor<1x2048xf32> loc(#loc5310)
    %27637 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27638 = stablehlo.reduce(%27636 init: %27637) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27639 = stablehlo.multiply %27635, %27638 : tensor<f32> loc(#loc5310)
    %27640 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27641 = stablehlo.broadcast_in_dim %27640, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27642 = stablehlo.broadcast_in_dim %27639, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27643 = stablehlo.concatenate %27641, %27642, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27644 = stablehlo.abs %25526 : tensor<1x8192x2048xf32> loc(#loc5307)
    %27645 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27646 = stablehlo.reduce(%27644 init: %27645) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27647 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27648 = stablehlo.maximum %27646, %27647 : tensor<f32> loc(#loc5309)
    %27649 = stablehlo.broadcast_in_dim %27648, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %27650 = stablehlo.divide %25526, %27649 : tensor<1x8192x2048xf32> loc(#loc5314)
    %27651 = stablehlo.multiply %27648, %27648 : tensor<f32> loc(#loc5310)
    %27652 = stablehlo.multiply %27650, %27650 : tensor<1x8192x2048xf32> loc(#loc5310)
    %27653 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27654 = stablehlo.reduce(%27652 init: %27653) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27655 = stablehlo.multiply %27651, %27654 : tensor<f32> loc(#loc5310)
    %27656 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27657 = stablehlo.broadcast_in_dim %27656, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27658 = stablehlo.broadcast_in_dim %27655, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27659 = stablehlo.concatenate %27657, %27658, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27660 = stablehlo.abs %25527 : tensor<1x2048xf32> loc(#loc5307)
    %27661 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27662 = stablehlo.reduce(%27660 init: %27661) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27663 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27664 = stablehlo.maximum %27662, %27663 : tensor<f32> loc(#loc5309)
    %27665 = stablehlo.broadcast_in_dim %27664, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27666 = stablehlo.divide %25527, %27665 : tensor<1x2048xf32> loc(#loc5314)
    %27667 = stablehlo.multiply %27664, %27664 : tensor<f32> loc(#loc5310)
    %27668 = stablehlo.multiply %27666, %27666 : tensor<1x2048xf32> loc(#loc5310)
    %27669 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27670 = stablehlo.reduce(%27668 init: %27669) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27671 = stablehlo.multiply %27667, %27670 : tensor<f32> loc(#loc5310)
    %27672 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27673 = stablehlo.broadcast_in_dim %27672, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27674 = stablehlo.broadcast_in_dim %27671, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27675 = stablehlo.concatenate %27673, %27674, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27676 = stablehlo.abs %25528 : tensor<1x2048xf32> loc(#loc5307)
    %27677 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27678 = stablehlo.reduce(%27676 init: %27677) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27679 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27680 = stablehlo.maximum %27678, %27679 : tensor<f32> loc(#loc5309)
    %27681 = stablehlo.broadcast_in_dim %27680, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27682 = stablehlo.divide %25528, %27681 : tensor<1x2048xf32> loc(#loc5314)
    %27683 = stablehlo.multiply %27680, %27680 : tensor<f32> loc(#loc5310)
    %27684 = stablehlo.multiply %27682, %27682 : tensor<1x2048xf32> loc(#loc5310)
    %27685 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27686 = stablehlo.reduce(%27684 init: %27685) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27687 = stablehlo.multiply %27683, %27686 : tensor<f32> loc(#loc5310)
    %27688 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27689 = stablehlo.broadcast_in_dim %27688, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27690 = stablehlo.broadcast_in_dim %27687, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27691 = stablehlo.concatenate %27689, %27690, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27692 = stablehlo.abs %25529 : tensor<1x2048xf32> loc(#loc5307)
    %27693 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27694 = stablehlo.reduce(%27692 init: %27693) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27695 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27696 = stablehlo.maximum %27694, %27695 : tensor<f32> loc(#loc5309)
    %27697 = stablehlo.broadcast_in_dim %27696, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27698 = stablehlo.divide %25529, %27697 : tensor<1x2048xf32> loc(#loc5314)
    %27699 = stablehlo.multiply %27696, %27696 : tensor<f32> loc(#loc5310)
    %27700 = stablehlo.multiply %27698, %27698 : tensor<1x2048xf32> loc(#loc5310)
    %27701 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27702 = stablehlo.reduce(%27700 init: %27701) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27703 = stablehlo.multiply %27699, %27702 : tensor<f32> loc(#loc5310)
    %27704 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27705 = stablehlo.broadcast_in_dim %27704, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27706 = stablehlo.broadcast_in_dim %27703, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27707 = stablehlo.concatenate %27705, %27706, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27708 = stablehlo.abs %25530 : tensor<1x2048xf32> loc(#loc5307)
    %27709 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27710 = stablehlo.reduce(%27708 init: %27709) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27711 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27712 = stablehlo.maximum %27710, %27711 : tensor<f32> loc(#loc5309)
    %27713 = stablehlo.broadcast_in_dim %27712, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27714 = stablehlo.divide %25530, %27713 : tensor<1x2048xf32> loc(#loc5314)
    %27715 = stablehlo.multiply %27712, %27712 : tensor<f32> loc(#loc5310)
    %27716 = stablehlo.multiply %27714, %27714 : tensor<1x2048xf32> loc(#loc5310)
    %27717 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27718 = stablehlo.reduce(%27716 init: %27717) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27719 = stablehlo.multiply %27715, %27718 : tensor<f32> loc(#loc5310)
    %27720 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27721 = stablehlo.broadcast_in_dim %27720, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27722 = stablehlo.broadcast_in_dim %27719, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27723 = stablehlo.concatenate %27721, %27722, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27724 = stablehlo.abs %25531 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %27725 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %27726 = stablehlo.reduce(%27724 init: %27725) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %27727 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27728 = stablehlo.maximum %27726, %27727 : tensor<f32> loc(#loc5309)
    %27729 = stablehlo.broadcast_in_dim %27728, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27730 = stablehlo.divide %25531, %27729 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27731 = stablehlo.multiply %27728, %27728 : tensor<f32> loc(#loc5310)
    %27732 = stablehlo.multiply %27730, %27730 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %27733 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %27734 = stablehlo.reduce(%27732 init: %27733) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %27735 = stablehlo.multiply %27731, %27734 : tensor<f32> loc(#loc5310)
    %27736 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %27737 = stablehlo.broadcast_in_dim %27736, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27738 = stablehlo.broadcast_in_dim %27735, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27739 = stablehlo.concatenate %27737, %27738, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27740 = stablehlo.abs %25532 : tensor<1x64xf32> loc(#loc5307)
    %27741 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27742 = stablehlo.reduce(%27740 init: %27741) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27743 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27744 = stablehlo.maximum %27742, %27743 : tensor<f32> loc(#loc5309)
    %27745 = stablehlo.broadcast_in_dim %27744, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %27746 = stablehlo.divide %25532, %27745 : tensor<1x64xf32> loc(#loc5314)
    %27747 = stablehlo.multiply %27744, %27744 : tensor<f32> loc(#loc5310)
    %27748 = stablehlo.multiply %27746, %27746 : tensor<1x64xf32> loc(#loc5310)
    %27749 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27750 = stablehlo.reduce(%27748 init: %27749) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27751 = stablehlo.multiply %27747, %27750 : tensor<f32> loc(#loc5310)
    %27752 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %27753 = stablehlo.broadcast_in_dim %27752, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27754 = stablehlo.broadcast_in_dim %27751, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27755 = stablehlo.concatenate %27753, %27754, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27756 = stablehlo.abs %25533 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %27757 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %27758 = stablehlo.reduce(%27756 init: %27757) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %27759 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27760 = stablehlo.maximum %27758, %27759 : tensor<f32> loc(#loc5309)
    %27761 = stablehlo.broadcast_in_dim %27760, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %27762 = stablehlo.divide %25533, %27761 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %27763 = stablehlo.multiply %27760, %27760 : tensor<f32> loc(#loc5310)
    %27764 = stablehlo.multiply %27762, %27762 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %27765 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %27766 = stablehlo.reduce(%27764 init: %27765) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %27767 = stablehlo.multiply %27763, %27766 : tensor<f32> loc(#loc5310)
    %27768 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %27769 = stablehlo.broadcast_in_dim %27768, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27770 = stablehlo.broadcast_in_dim %27767, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27771 = stablehlo.concatenate %27769, %27770, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27772 = stablehlo.abs %25534 : tensor<1x8192xf32> loc(#loc5307)
    %27773 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27774 = stablehlo.reduce(%27772 init: %27773) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27775 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27776 = stablehlo.maximum %27774, %27775 : tensor<f32> loc(#loc5309)
    %27777 = stablehlo.broadcast_in_dim %27776, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %27778 = stablehlo.divide %25534, %27777 : tensor<1x8192xf32> loc(#loc5314)
    %27779 = stablehlo.multiply %27776, %27776 : tensor<f32> loc(#loc5310)
    %27780 = stablehlo.multiply %27778, %27778 : tensor<1x8192xf32> loc(#loc5310)
    %27781 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27782 = stablehlo.reduce(%27780 init: %27781) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27783 = stablehlo.multiply %27779, %27782 : tensor<f32> loc(#loc5310)
    %27784 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %27785 = stablehlo.broadcast_in_dim %27784, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27786 = stablehlo.broadcast_in_dim %27783, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27787 = stablehlo.concatenate %27785, %27786, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27788 = stablehlo.abs %25535 : tensor<1x2048x8192xf32> loc(#loc5307)
    %27789 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27790 = stablehlo.reduce(%27788 init: %27789) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27791 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27792 = stablehlo.maximum %27790, %27791 : tensor<f32> loc(#loc5309)
    %27793 = stablehlo.broadcast_in_dim %27792, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %27794 = stablehlo.divide %25535, %27793 : tensor<1x2048x8192xf32> loc(#loc5314)
    %27795 = stablehlo.multiply %27792, %27792 : tensor<f32> loc(#loc5310)
    %27796 = stablehlo.multiply %27794, %27794 : tensor<1x2048x8192xf32> loc(#loc5310)
    %27797 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27798 = stablehlo.reduce(%27796 init: %27797) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27799 = stablehlo.multiply %27795, %27798 : tensor<f32> loc(#loc5310)
    %27800 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27801 = stablehlo.broadcast_in_dim %27800, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27802 = stablehlo.broadcast_in_dim %27799, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27803 = stablehlo.concatenate %27801, %27802, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27804 = stablehlo.abs %25536 : tensor<1x2048xf32> loc(#loc5307)
    %27805 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27806 = stablehlo.reduce(%27804 init: %27805) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27807 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27808 = stablehlo.maximum %27806, %27807 : tensor<f32> loc(#loc5309)
    %27809 = stablehlo.broadcast_in_dim %27808, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27810 = stablehlo.divide %25536, %27809 : tensor<1x2048xf32> loc(#loc5314)
    %27811 = stablehlo.multiply %27808, %27808 : tensor<f32> loc(#loc5310)
    %27812 = stablehlo.multiply %27810, %27810 : tensor<1x2048xf32> loc(#loc5310)
    %27813 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27814 = stablehlo.reduce(%27812 init: %27813) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27815 = stablehlo.multiply %27811, %27814 : tensor<f32> loc(#loc5310)
    %27816 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27817 = stablehlo.broadcast_in_dim %27816, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27818 = stablehlo.broadcast_in_dim %27815, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27819 = stablehlo.concatenate %27817, %27818, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27820 = stablehlo.abs %25537 : tensor<1x8192x2048xf32> loc(#loc5307)
    %27821 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27822 = stablehlo.reduce(%27820 init: %27821) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27823 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27824 = stablehlo.maximum %27822, %27823 : tensor<f32> loc(#loc5309)
    %27825 = stablehlo.broadcast_in_dim %27824, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %27826 = stablehlo.divide %25537, %27825 : tensor<1x8192x2048xf32> loc(#loc5314)
    %27827 = stablehlo.multiply %27824, %27824 : tensor<f32> loc(#loc5310)
    %27828 = stablehlo.multiply %27826, %27826 : tensor<1x8192x2048xf32> loc(#loc5310)
    %27829 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27830 = stablehlo.reduce(%27828 init: %27829) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27831 = stablehlo.multiply %27827, %27830 : tensor<f32> loc(#loc5310)
    %27832 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27833 = stablehlo.broadcast_in_dim %27832, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27834 = stablehlo.broadcast_in_dim %27831, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27835 = stablehlo.concatenate %27833, %27834, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27836 = stablehlo.abs %25538 : tensor<1x2048xf32> loc(#loc5307)
    %27837 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27838 = stablehlo.reduce(%27836 init: %27837) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27839 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27840 = stablehlo.maximum %27838, %27839 : tensor<f32> loc(#loc5309)
    %27841 = stablehlo.broadcast_in_dim %27840, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27842 = stablehlo.divide %25538, %27841 : tensor<1x2048xf32> loc(#loc5314)
    %27843 = stablehlo.multiply %27840, %27840 : tensor<f32> loc(#loc5310)
    %27844 = stablehlo.multiply %27842, %27842 : tensor<1x2048xf32> loc(#loc5310)
    %27845 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27846 = stablehlo.reduce(%27844 init: %27845) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27847 = stablehlo.multiply %27843, %27846 : tensor<f32> loc(#loc5310)
    %27848 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27849 = stablehlo.broadcast_in_dim %27848, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27850 = stablehlo.broadcast_in_dim %27847, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27851 = stablehlo.concatenate %27849, %27850, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27852 = stablehlo.abs %25539 : tensor<1x2048xf32> loc(#loc5307)
    %27853 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27854 = stablehlo.reduce(%27852 init: %27853) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27855 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27856 = stablehlo.maximum %27854, %27855 : tensor<f32> loc(#loc5309)
    %27857 = stablehlo.broadcast_in_dim %27856, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27858 = stablehlo.divide %25539, %27857 : tensor<1x2048xf32> loc(#loc5314)
    %27859 = stablehlo.multiply %27856, %27856 : tensor<f32> loc(#loc5310)
    %27860 = stablehlo.multiply %27858, %27858 : tensor<1x2048xf32> loc(#loc5310)
    %27861 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27862 = stablehlo.reduce(%27860 init: %27861) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27863 = stablehlo.multiply %27859, %27862 : tensor<f32> loc(#loc5310)
    %27864 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27865 = stablehlo.broadcast_in_dim %27864, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27866 = stablehlo.broadcast_in_dim %27863, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27867 = stablehlo.concatenate %27865, %27866, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27868 = stablehlo.abs %25540 : tensor<1x2048xf32> loc(#loc5307)
    %27869 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27870 = stablehlo.reduce(%27868 init: %27869) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27871 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27872 = stablehlo.maximum %27870, %27871 : tensor<f32> loc(#loc5309)
    %27873 = stablehlo.broadcast_in_dim %27872, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27874 = stablehlo.divide %25540, %27873 : tensor<1x2048xf32> loc(#loc5314)
    %27875 = stablehlo.multiply %27872, %27872 : tensor<f32> loc(#loc5310)
    %27876 = stablehlo.multiply %27874, %27874 : tensor<1x2048xf32> loc(#loc5310)
    %27877 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27878 = stablehlo.reduce(%27876 init: %27877) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27879 = stablehlo.multiply %27875, %27878 : tensor<f32> loc(#loc5310)
    %27880 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27881 = stablehlo.broadcast_in_dim %27880, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27882 = stablehlo.broadcast_in_dim %27879, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27883 = stablehlo.concatenate %27881, %27882, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27884 = stablehlo.abs %25541 : tensor<1x2048xf32> loc(#loc5307)
    %27885 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27886 = stablehlo.reduce(%27884 init: %27885) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27887 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27888 = stablehlo.maximum %27886, %27887 : tensor<f32> loc(#loc5309)
    %27889 = stablehlo.broadcast_in_dim %27888, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27890 = stablehlo.divide %25541, %27889 : tensor<1x2048xf32> loc(#loc5314)
    %27891 = stablehlo.multiply %27888, %27888 : tensor<f32> loc(#loc5310)
    %27892 = stablehlo.multiply %27890, %27890 : tensor<1x2048xf32> loc(#loc5310)
    %27893 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27894 = stablehlo.reduce(%27892 init: %27893) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27895 = stablehlo.multiply %27891, %27894 : tensor<f32> loc(#loc5310)
    %27896 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27897 = stablehlo.broadcast_in_dim %27896, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27898 = stablehlo.broadcast_in_dim %27895, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27899 = stablehlo.concatenate %27897, %27898, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27900 = stablehlo.abs %25542 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %27901 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %27902 = stablehlo.reduce(%27900 init: %27901) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %27903 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27904 = stablehlo.maximum %27902, %27903 : tensor<f32> loc(#loc5309)
    %27905 = stablehlo.broadcast_in_dim %27904, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27906 = stablehlo.divide %25542, %27905 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %27907 = stablehlo.multiply %27904, %27904 : tensor<f32> loc(#loc5310)
    %27908 = stablehlo.multiply %27906, %27906 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %27909 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %27910 = stablehlo.reduce(%27908 init: %27909) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %27911 = stablehlo.multiply %27907, %27910 : tensor<f32> loc(#loc5310)
    %27912 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %27913 = stablehlo.broadcast_in_dim %27912, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27914 = stablehlo.broadcast_in_dim %27911, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27915 = stablehlo.concatenate %27913, %27914, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27916 = stablehlo.abs %25543 : tensor<1x64xf32> loc(#loc5307)
    %27917 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27918 = stablehlo.reduce(%27916 init: %27917) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27919 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27920 = stablehlo.maximum %27918, %27919 : tensor<f32> loc(#loc5309)
    %27921 = stablehlo.broadcast_in_dim %27920, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %27922 = stablehlo.divide %25543, %27921 : tensor<1x64xf32> loc(#loc5314)
    %27923 = stablehlo.multiply %27920, %27920 : tensor<f32> loc(#loc5310)
    %27924 = stablehlo.multiply %27922, %27922 : tensor<1x64xf32> loc(#loc5310)
    %27925 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27926 = stablehlo.reduce(%27924 init: %27925) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27927 = stablehlo.multiply %27923, %27926 : tensor<f32> loc(#loc5310)
    %27928 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %27929 = stablehlo.broadcast_in_dim %27928, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27930 = stablehlo.broadcast_in_dim %27927, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27931 = stablehlo.concatenate %27929, %27930, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27932 = stablehlo.abs %25544 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %27933 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %27934 = stablehlo.reduce(%27932 init: %27933) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %27935 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27936 = stablehlo.maximum %27934, %27935 : tensor<f32> loc(#loc5309)
    %27937 = stablehlo.broadcast_in_dim %27936, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %27938 = stablehlo.divide %25544, %27937 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %27939 = stablehlo.multiply %27936, %27936 : tensor<f32> loc(#loc5310)
    %27940 = stablehlo.multiply %27938, %27938 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %27941 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %27942 = stablehlo.reduce(%27940 init: %27941) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %27943 = stablehlo.multiply %27939, %27942 : tensor<f32> loc(#loc5310)
    %27944 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %27945 = stablehlo.broadcast_in_dim %27944, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27946 = stablehlo.broadcast_in_dim %27943, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27947 = stablehlo.concatenate %27945, %27946, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27948 = stablehlo.abs %25545 : tensor<1x8192xf32> loc(#loc5307)
    %27949 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27950 = stablehlo.reduce(%27948 init: %27949) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27951 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27952 = stablehlo.maximum %27950, %27951 : tensor<f32> loc(#loc5309)
    %27953 = stablehlo.broadcast_in_dim %27952, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %27954 = stablehlo.divide %25545, %27953 : tensor<1x8192xf32> loc(#loc5314)
    %27955 = stablehlo.multiply %27952, %27952 : tensor<f32> loc(#loc5310)
    %27956 = stablehlo.multiply %27954, %27954 : tensor<1x8192xf32> loc(#loc5310)
    %27957 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %27958 = stablehlo.reduce(%27956 init: %27957) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %27959 = stablehlo.multiply %27955, %27958 : tensor<f32> loc(#loc5310)
    %27960 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %27961 = stablehlo.broadcast_in_dim %27960, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27962 = stablehlo.broadcast_in_dim %27959, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27963 = stablehlo.concatenate %27961, %27962, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27964 = stablehlo.abs %25546 : tensor<1x2048x8192xf32> loc(#loc5307)
    %27965 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27966 = stablehlo.reduce(%27964 init: %27965) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27967 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27968 = stablehlo.maximum %27966, %27967 : tensor<f32> loc(#loc5309)
    %27969 = stablehlo.broadcast_in_dim %27968, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %27970 = stablehlo.divide %25546, %27969 : tensor<1x2048x8192xf32> loc(#loc5314)
    %27971 = stablehlo.multiply %27968, %27968 : tensor<f32> loc(#loc5310)
    %27972 = stablehlo.multiply %27970, %27970 : tensor<1x2048x8192xf32> loc(#loc5310)
    %27973 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %27974 = stablehlo.reduce(%27972 init: %27973) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %27975 = stablehlo.multiply %27971, %27974 : tensor<f32> loc(#loc5310)
    %27976 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %27977 = stablehlo.broadcast_in_dim %27976, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27978 = stablehlo.broadcast_in_dim %27975, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27979 = stablehlo.concatenate %27977, %27978, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27980 = stablehlo.abs %25547 : tensor<1x2048xf32> loc(#loc5307)
    %27981 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %27982 = stablehlo.reduce(%27980 init: %27981) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %27983 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %27984 = stablehlo.maximum %27982, %27983 : tensor<f32> loc(#loc5309)
    %27985 = stablehlo.broadcast_in_dim %27984, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %27986 = stablehlo.divide %25547, %27985 : tensor<1x2048xf32> loc(#loc5314)
    %27987 = stablehlo.multiply %27984, %27984 : tensor<f32> loc(#loc5310)
    %27988 = stablehlo.multiply %27986, %27986 : tensor<1x2048xf32> loc(#loc5310)
    %27989 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %27990 = stablehlo.reduce(%27988 init: %27989) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %27991 = stablehlo.multiply %27987, %27990 : tensor<f32> loc(#loc5310)
    %27992 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %27993 = stablehlo.broadcast_in_dim %27992, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27994 = stablehlo.broadcast_in_dim %27991, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %27995 = stablehlo.concatenate %27993, %27994, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %27996 = stablehlo.abs %25548 : tensor<1x8192x2048xf32> loc(#loc5307)
    %27997 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %27998 = stablehlo.reduce(%27996 init: %27997) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %27999 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28000 = stablehlo.maximum %27998, %27999 : tensor<f32> loc(#loc5309)
    %28001 = stablehlo.broadcast_in_dim %28000, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %28002 = stablehlo.divide %25548, %28001 : tensor<1x8192x2048xf32> loc(#loc5314)
    %28003 = stablehlo.multiply %28000, %28000 : tensor<f32> loc(#loc5310)
    %28004 = stablehlo.multiply %28002, %28002 : tensor<1x8192x2048xf32> loc(#loc5310)
    %28005 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %28006 = stablehlo.reduce(%28004 init: %28005) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %28007 = stablehlo.multiply %28003, %28006 : tensor<f32> loc(#loc5310)
    %28008 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %28009 = stablehlo.broadcast_in_dim %28008, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28010 = stablehlo.broadcast_in_dim %28007, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28011 = stablehlo.concatenate %28009, %28010, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28012 = stablehlo.abs %25549 : tensor<1x2048xf32> loc(#loc5307)
    %28013 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28014 = stablehlo.reduce(%28012 init: %28013) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28015 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28016 = stablehlo.maximum %28014, %28015 : tensor<f32> loc(#loc5309)
    %28017 = stablehlo.broadcast_in_dim %28016, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28018 = stablehlo.divide %25549, %28017 : tensor<1x2048xf32> loc(#loc5314)
    %28019 = stablehlo.multiply %28016, %28016 : tensor<f32> loc(#loc5310)
    %28020 = stablehlo.multiply %28018, %28018 : tensor<1x2048xf32> loc(#loc5310)
    %28021 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28022 = stablehlo.reduce(%28020 init: %28021) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28023 = stablehlo.multiply %28019, %28022 : tensor<f32> loc(#loc5310)
    %28024 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28025 = stablehlo.broadcast_in_dim %28024, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28026 = stablehlo.broadcast_in_dim %28023, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28027 = stablehlo.concatenate %28025, %28026, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28028 = stablehlo.abs %25550 : tensor<1x2048xf32> loc(#loc5307)
    %28029 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28030 = stablehlo.reduce(%28028 init: %28029) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28031 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28032 = stablehlo.maximum %28030, %28031 : tensor<f32> loc(#loc5309)
    %28033 = stablehlo.broadcast_in_dim %28032, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28034 = stablehlo.divide %25550, %28033 : tensor<1x2048xf32> loc(#loc5314)
    %28035 = stablehlo.multiply %28032, %28032 : tensor<f32> loc(#loc5310)
    %28036 = stablehlo.multiply %28034, %28034 : tensor<1x2048xf32> loc(#loc5310)
    %28037 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28038 = stablehlo.reduce(%28036 init: %28037) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28039 = stablehlo.multiply %28035, %28038 : tensor<f32> loc(#loc5310)
    %28040 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28041 = stablehlo.broadcast_in_dim %28040, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28042 = stablehlo.broadcast_in_dim %28039, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28043 = stablehlo.concatenate %28041, %28042, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28044 = stablehlo.abs %25551 : tensor<1x2048xf32> loc(#loc5307)
    %28045 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28046 = stablehlo.reduce(%28044 init: %28045) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28047 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28048 = stablehlo.maximum %28046, %28047 : tensor<f32> loc(#loc5309)
    %28049 = stablehlo.broadcast_in_dim %28048, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28050 = stablehlo.divide %25551, %28049 : tensor<1x2048xf32> loc(#loc5314)
    %28051 = stablehlo.multiply %28048, %28048 : tensor<f32> loc(#loc5310)
    %28052 = stablehlo.multiply %28050, %28050 : tensor<1x2048xf32> loc(#loc5310)
    %28053 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28054 = stablehlo.reduce(%28052 init: %28053) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28055 = stablehlo.multiply %28051, %28054 : tensor<f32> loc(#loc5310)
    %28056 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28057 = stablehlo.broadcast_in_dim %28056, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28058 = stablehlo.broadcast_in_dim %28055, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28059 = stablehlo.concatenate %28057, %28058, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28060 = stablehlo.abs %25552 : tensor<1x2048xf32> loc(#loc5307)
    %28061 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28062 = stablehlo.reduce(%28060 init: %28061) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28063 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28064 = stablehlo.maximum %28062, %28063 : tensor<f32> loc(#loc5309)
    %28065 = stablehlo.broadcast_in_dim %28064, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28066 = stablehlo.divide %25552, %28065 : tensor<1x2048xf32> loc(#loc5314)
    %28067 = stablehlo.multiply %28064, %28064 : tensor<f32> loc(#loc5310)
    %28068 = stablehlo.multiply %28066, %28066 : tensor<1x2048xf32> loc(#loc5310)
    %28069 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28070 = stablehlo.reduce(%28068 init: %28069) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28071 = stablehlo.multiply %28067, %28070 : tensor<f32> loc(#loc5310)
    %28072 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28073 = stablehlo.broadcast_in_dim %28072, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28074 = stablehlo.broadcast_in_dim %28071, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28075 = stablehlo.concatenate %28073, %28074, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28076 = stablehlo.abs %25553 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %28077 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %28078 = stablehlo.reduce(%28076 init: %28077) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %28079 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28080 = stablehlo.maximum %28078, %28079 : tensor<f32> loc(#loc5309)
    %28081 = stablehlo.broadcast_in_dim %28080, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %28082 = stablehlo.divide %25553, %28081 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %28083 = stablehlo.multiply %28080, %28080 : tensor<f32> loc(#loc5310)
    %28084 = stablehlo.multiply %28082, %28082 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %28085 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %28086 = stablehlo.reduce(%28084 init: %28085) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %28087 = stablehlo.multiply %28083, %28086 : tensor<f32> loc(#loc5310)
    %28088 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %28089 = stablehlo.broadcast_in_dim %28088, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28090 = stablehlo.broadcast_in_dim %28087, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28091 = stablehlo.concatenate %28089, %28090, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28092 = stablehlo.abs %25554 : tensor<1x64xf32> loc(#loc5307)
    %28093 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28094 = stablehlo.reduce(%28092 init: %28093) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28095 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28096 = stablehlo.maximum %28094, %28095 : tensor<f32> loc(#loc5309)
    %28097 = stablehlo.broadcast_in_dim %28096, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %28098 = stablehlo.divide %25554, %28097 : tensor<1x64xf32> loc(#loc5314)
    %28099 = stablehlo.multiply %28096, %28096 : tensor<f32> loc(#loc5310)
    %28100 = stablehlo.multiply %28098, %28098 : tensor<1x64xf32> loc(#loc5310)
    %28101 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %28102 = stablehlo.reduce(%28100 init: %28101) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %28103 = stablehlo.multiply %28099, %28102 : tensor<f32> loc(#loc5310)
    %28104 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %28105 = stablehlo.broadcast_in_dim %28104, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28106 = stablehlo.broadcast_in_dim %28103, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28107 = stablehlo.concatenate %28105, %28106, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28108 = stablehlo.abs %25555 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %28109 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %28110 = stablehlo.reduce(%28108 init: %28109) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %28111 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28112 = stablehlo.maximum %28110, %28111 : tensor<f32> loc(#loc5309)
    %28113 = stablehlo.broadcast_in_dim %28112, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %28114 = stablehlo.divide %25555, %28113 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %28115 = stablehlo.multiply %28112, %28112 : tensor<f32> loc(#loc5310)
    %28116 = stablehlo.multiply %28114, %28114 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %28117 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %28118 = stablehlo.reduce(%28116 init: %28117) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %28119 = stablehlo.multiply %28115, %28118 : tensor<f32> loc(#loc5310)
    %28120 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %28121 = stablehlo.broadcast_in_dim %28120, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28122 = stablehlo.broadcast_in_dim %28119, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28123 = stablehlo.concatenate %28121, %28122, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28124 = stablehlo.abs %25556 : tensor<1x8192xf32> loc(#loc5307)
    %28125 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28126 = stablehlo.reduce(%28124 init: %28125) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28127 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28128 = stablehlo.maximum %28126, %28127 : tensor<f32> loc(#loc5309)
    %28129 = stablehlo.broadcast_in_dim %28128, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %28130 = stablehlo.divide %25556, %28129 : tensor<1x8192xf32> loc(#loc5314)
    %28131 = stablehlo.multiply %28128, %28128 : tensor<f32> loc(#loc5310)
    %28132 = stablehlo.multiply %28130, %28130 : tensor<1x8192xf32> loc(#loc5310)
    %28133 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %28134 = stablehlo.reduce(%28132 init: %28133) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %28135 = stablehlo.multiply %28131, %28134 : tensor<f32> loc(#loc5310)
    %28136 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %28137 = stablehlo.broadcast_in_dim %28136, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28138 = stablehlo.broadcast_in_dim %28135, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28139 = stablehlo.concatenate %28137, %28138, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28140 = stablehlo.abs %25557 : tensor<1x2048x8192xf32> loc(#loc5307)
    %28141 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %28142 = stablehlo.reduce(%28140 init: %28141) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %28143 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28144 = stablehlo.maximum %28142, %28143 : tensor<f32> loc(#loc5309)
    %28145 = stablehlo.broadcast_in_dim %28144, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %28146 = stablehlo.divide %25557, %28145 : tensor<1x2048x8192xf32> loc(#loc5314)
    %28147 = stablehlo.multiply %28144, %28144 : tensor<f32> loc(#loc5310)
    %28148 = stablehlo.multiply %28146, %28146 : tensor<1x2048x8192xf32> loc(#loc5310)
    %28149 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %28150 = stablehlo.reduce(%28148 init: %28149) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %28151 = stablehlo.multiply %28147, %28150 : tensor<f32> loc(#loc5310)
    %28152 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %28153 = stablehlo.broadcast_in_dim %28152, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28154 = stablehlo.broadcast_in_dim %28151, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28155 = stablehlo.concatenate %28153, %28154, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28156 = stablehlo.abs %25558 : tensor<1x2048xf32> loc(#loc5307)
    %28157 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28158 = stablehlo.reduce(%28156 init: %28157) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28159 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28160 = stablehlo.maximum %28158, %28159 : tensor<f32> loc(#loc5309)
    %28161 = stablehlo.broadcast_in_dim %28160, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28162 = stablehlo.divide %25558, %28161 : tensor<1x2048xf32> loc(#loc5314)
    %28163 = stablehlo.multiply %28160, %28160 : tensor<f32> loc(#loc5310)
    %28164 = stablehlo.multiply %28162, %28162 : tensor<1x2048xf32> loc(#loc5310)
    %28165 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28166 = stablehlo.reduce(%28164 init: %28165) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28167 = stablehlo.multiply %28163, %28166 : tensor<f32> loc(#loc5310)
    %28168 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28169 = stablehlo.broadcast_in_dim %28168, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28170 = stablehlo.broadcast_in_dim %28167, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28171 = stablehlo.concatenate %28169, %28170, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28172 = stablehlo.abs %25559 : tensor<1x8192x2048xf32> loc(#loc5307)
    %28173 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %28174 = stablehlo.reduce(%28172 init: %28173) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %28175 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28176 = stablehlo.maximum %28174, %28175 : tensor<f32> loc(#loc5309)
    %28177 = stablehlo.broadcast_in_dim %28176, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %28178 = stablehlo.divide %25559, %28177 : tensor<1x8192x2048xf32> loc(#loc5314)
    %28179 = stablehlo.multiply %28176, %28176 : tensor<f32> loc(#loc5310)
    %28180 = stablehlo.multiply %28178, %28178 : tensor<1x8192x2048xf32> loc(#loc5310)
    %28181 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %28182 = stablehlo.reduce(%28180 init: %28181) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %28183 = stablehlo.multiply %28179, %28182 : tensor<f32> loc(#loc5310)
    %28184 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %28185 = stablehlo.broadcast_in_dim %28184, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28186 = stablehlo.broadcast_in_dim %28183, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28187 = stablehlo.concatenate %28185, %28186, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28188 = stablehlo.abs %25560 : tensor<1x2048xf32> loc(#loc5307)
    %28189 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28190 = stablehlo.reduce(%28188 init: %28189) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28191 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28192 = stablehlo.maximum %28190, %28191 : tensor<f32> loc(#loc5309)
    %28193 = stablehlo.broadcast_in_dim %28192, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28194 = stablehlo.divide %25560, %28193 : tensor<1x2048xf32> loc(#loc5314)
    %28195 = stablehlo.multiply %28192, %28192 : tensor<f32> loc(#loc5310)
    %28196 = stablehlo.multiply %28194, %28194 : tensor<1x2048xf32> loc(#loc5310)
    %28197 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28198 = stablehlo.reduce(%28196 init: %28197) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28199 = stablehlo.multiply %28195, %28198 : tensor<f32> loc(#loc5310)
    %28200 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28201 = stablehlo.broadcast_in_dim %28200, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28202 = stablehlo.broadcast_in_dim %28199, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28203 = stablehlo.concatenate %28201, %28202, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28204 = stablehlo.abs %25561 : tensor<1x2048xf32> loc(#loc5307)
    %28205 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28206 = stablehlo.reduce(%28204 init: %28205) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28207 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28208 = stablehlo.maximum %28206, %28207 : tensor<f32> loc(#loc5309)
    %28209 = stablehlo.broadcast_in_dim %28208, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28210 = stablehlo.divide %25561, %28209 : tensor<1x2048xf32> loc(#loc5314)
    %28211 = stablehlo.multiply %28208, %28208 : tensor<f32> loc(#loc5310)
    %28212 = stablehlo.multiply %28210, %28210 : tensor<1x2048xf32> loc(#loc5310)
    %28213 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28214 = stablehlo.reduce(%28212 init: %28213) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28215 = stablehlo.multiply %28211, %28214 : tensor<f32> loc(#loc5310)
    %28216 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28217 = stablehlo.broadcast_in_dim %28216, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28218 = stablehlo.broadcast_in_dim %28215, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28219 = stablehlo.concatenate %28217, %28218, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28220 = stablehlo.abs %25562 : tensor<1x2048xf32> loc(#loc5307)
    %28221 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28222 = stablehlo.reduce(%28220 init: %28221) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28223 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28224 = stablehlo.maximum %28222, %28223 : tensor<f32> loc(#loc5309)
    %28225 = stablehlo.broadcast_in_dim %28224, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28226 = stablehlo.divide %25562, %28225 : tensor<1x2048xf32> loc(#loc5314)
    %28227 = stablehlo.multiply %28224, %28224 : tensor<f32> loc(#loc5310)
    %28228 = stablehlo.multiply %28226, %28226 : tensor<1x2048xf32> loc(#loc5310)
    %28229 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28230 = stablehlo.reduce(%28228 init: %28229) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28231 = stablehlo.multiply %28227, %28230 : tensor<f32> loc(#loc5310)
    %28232 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28233 = stablehlo.broadcast_in_dim %28232, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28234 = stablehlo.broadcast_in_dim %28231, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28235 = stablehlo.concatenate %28233, %28234, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28236 = stablehlo.abs %25563 : tensor<1x2048xf32> loc(#loc5307)
    %28237 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28238 = stablehlo.reduce(%28236 init: %28237) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28239 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28240 = stablehlo.maximum %28238, %28239 : tensor<f32> loc(#loc5309)
    %28241 = stablehlo.broadcast_in_dim %28240, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28242 = stablehlo.divide %25563, %28241 : tensor<1x2048xf32> loc(#loc5314)
    %28243 = stablehlo.multiply %28240, %28240 : tensor<f32> loc(#loc5310)
    %28244 = stablehlo.multiply %28242, %28242 : tensor<1x2048xf32> loc(#loc5310)
    %28245 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28246 = stablehlo.reduce(%28244 init: %28245) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28247 = stablehlo.multiply %28243, %28246 : tensor<f32> loc(#loc5310)
    %28248 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28249 = stablehlo.broadcast_in_dim %28248, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28250 = stablehlo.broadcast_in_dim %28247, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28251 = stablehlo.concatenate %28249, %28250, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28252 = stablehlo.abs %25564 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %28253 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %28254 = stablehlo.reduce(%28252 init: %28253) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %28255 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28256 = stablehlo.maximum %28254, %28255 : tensor<f32> loc(#loc5309)
    %28257 = stablehlo.broadcast_in_dim %28256, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %28258 = stablehlo.divide %25564, %28257 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %28259 = stablehlo.multiply %28256, %28256 : tensor<f32> loc(#loc5310)
    %28260 = stablehlo.multiply %28258, %28258 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %28261 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %28262 = stablehlo.reduce(%28260 init: %28261) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %28263 = stablehlo.multiply %28259, %28262 : tensor<f32> loc(#loc5310)
    %28264 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %28265 = stablehlo.broadcast_in_dim %28264, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28266 = stablehlo.broadcast_in_dim %28263, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28267 = stablehlo.concatenate %28265, %28266, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28268 = stablehlo.abs %25565 : tensor<1x64xf32> loc(#loc5307)
    %28269 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28270 = stablehlo.reduce(%28268 init: %28269) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28271 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28272 = stablehlo.maximum %28270, %28271 : tensor<f32> loc(#loc5309)
    %28273 = stablehlo.broadcast_in_dim %28272, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %28274 = stablehlo.divide %25565, %28273 : tensor<1x64xf32> loc(#loc5314)
    %28275 = stablehlo.multiply %28272, %28272 : tensor<f32> loc(#loc5310)
    %28276 = stablehlo.multiply %28274, %28274 : tensor<1x64xf32> loc(#loc5310)
    %28277 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %28278 = stablehlo.reduce(%28276 init: %28277) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %28279 = stablehlo.multiply %28275, %28278 : tensor<f32> loc(#loc5310)
    %28280 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %28281 = stablehlo.broadcast_in_dim %28280, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28282 = stablehlo.broadcast_in_dim %28279, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28283 = stablehlo.concatenate %28281, %28282, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28284 = stablehlo.abs %25566 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %28285 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %28286 = stablehlo.reduce(%28284 init: %28285) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %28287 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28288 = stablehlo.maximum %28286, %28287 : tensor<f32> loc(#loc5309)
    %28289 = stablehlo.broadcast_in_dim %28288, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %28290 = stablehlo.divide %25566, %28289 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %28291 = stablehlo.multiply %28288, %28288 : tensor<f32> loc(#loc5310)
    %28292 = stablehlo.multiply %28290, %28290 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %28293 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %28294 = stablehlo.reduce(%28292 init: %28293) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %28295 = stablehlo.multiply %28291, %28294 : tensor<f32> loc(#loc5310)
    %28296 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %28297 = stablehlo.broadcast_in_dim %28296, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28298 = stablehlo.broadcast_in_dim %28295, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28299 = stablehlo.concatenate %28297, %28298, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28300 = stablehlo.abs %25567 : tensor<1x8192xf32> loc(#loc5307)
    %28301 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28302 = stablehlo.reduce(%28300 init: %28301) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28303 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28304 = stablehlo.maximum %28302, %28303 : tensor<f32> loc(#loc5309)
    %28305 = stablehlo.broadcast_in_dim %28304, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc5314)
    %28306 = stablehlo.divide %25567, %28305 : tensor<1x8192xf32> loc(#loc5314)
    %28307 = stablehlo.multiply %28304, %28304 : tensor<f32> loc(#loc5310)
    %28308 = stablehlo.multiply %28306, %28306 : tensor<1x8192xf32> loc(#loc5310)
    %28309 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %28310 = stablehlo.reduce(%28308 init: %28309) across dimensions = [0, 1] : (tensor<1x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %28311 = stablehlo.multiply %28307, %28310 : tensor<f32> loc(#loc5310)
    %28312 = stablehlo.constant dense<8.192000e+03> : tensor<f32> loc(#loc)
    %28313 = stablehlo.broadcast_in_dim %28312, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28314 = stablehlo.broadcast_in_dim %28311, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28315 = stablehlo.concatenate %28313, %28314, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28316 = stablehlo.abs %25568 : tensor<1x2048x8192xf32> loc(#loc5307)
    %28317 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %28318 = stablehlo.reduce(%28316 init: %28317) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %28319 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28320 = stablehlo.maximum %28318, %28319 : tensor<f32> loc(#loc5309)
    %28321 = stablehlo.broadcast_in_dim %28320, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc5314)
    %28322 = stablehlo.divide %25568, %28321 : tensor<1x2048x8192xf32> loc(#loc5314)
    %28323 = stablehlo.multiply %28320, %28320 : tensor<f32> loc(#loc5310)
    %28324 = stablehlo.multiply %28322, %28322 : tensor<1x2048x8192xf32> loc(#loc5310)
    %28325 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %28326 = stablehlo.reduce(%28324 init: %28325) across dimensions = [0, 1, 2] : (tensor<1x2048x8192xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %28327 = stablehlo.multiply %28323, %28326 : tensor<f32> loc(#loc5310)
    %28328 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %28329 = stablehlo.broadcast_in_dim %28328, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28330 = stablehlo.broadcast_in_dim %28327, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28331 = stablehlo.concatenate %28329, %28330, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28332 = stablehlo.abs %25569 : tensor<1x2048xf32> loc(#loc5307)
    %28333 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28334 = stablehlo.reduce(%28332 init: %28333) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28335 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28336 = stablehlo.maximum %28334, %28335 : tensor<f32> loc(#loc5309)
    %28337 = stablehlo.broadcast_in_dim %28336, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28338 = stablehlo.divide %25569, %28337 : tensor<1x2048xf32> loc(#loc5314)
    %28339 = stablehlo.multiply %28336, %28336 : tensor<f32> loc(#loc5310)
    %28340 = stablehlo.multiply %28338, %28338 : tensor<1x2048xf32> loc(#loc5310)
    %28341 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28342 = stablehlo.reduce(%28340 init: %28341) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28343 = stablehlo.multiply %28339, %28342 : tensor<f32> loc(#loc5310)
    %28344 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28345 = stablehlo.broadcast_in_dim %28344, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28346 = stablehlo.broadcast_in_dim %28343, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28347 = stablehlo.concatenate %28345, %28346, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28348 = stablehlo.abs %25570 : tensor<1x8192x2048xf32> loc(#loc5307)
    %28349 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5315)
    %28350 = stablehlo.reduce(%28348 init: %28349) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5315)
      stablehlo.return %29929 : tensor<f32> loc(#loc5315)
    } loc(#loc5315)
    %28351 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28352 = stablehlo.maximum %28350, %28351 : tensor<f32> loc(#loc5309)
    %28353 = stablehlo.broadcast_in_dim %28352, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc5314)
    %28354 = stablehlo.divide %25570, %28353 : tensor<1x8192x2048xf32> loc(#loc5314)
    %28355 = stablehlo.multiply %28352, %28352 : tensor<f32> loc(#loc5310)
    %28356 = stablehlo.multiply %28354, %28354 : tensor<1x8192x2048xf32> loc(#loc5310)
    %28357 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5241)
    %28358 = stablehlo.reduce(%28356 init: %28357) across dimensions = [0, 1, 2] : (tensor<1x8192x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5241)
      stablehlo.return %29929 : tensor<f32> loc(#loc5241)
    } loc(#loc5241)
    %28359 = stablehlo.multiply %28355, %28358 : tensor<f32> loc(#loc5310)
    %28360 = stablehlo.constant dense<0x4B800000> : tensor<f32> loc(#loc)
    %28361 = stablehlo.broadcast_in_dim %28360, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28362 = stablehlo.broadcast_in_dim %28359, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28363 = stablehlo.concatenate %28361, %28362, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28364 = stablehlo.abs %25571 : tensor<1x2048xf32> loc(#loc5307)
    %28365 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28366 = stablehlo.reduce(%28364 init: %28365) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28367 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28368 = stablehlo.maximum %28366, %28367 : tensor<f32> loc(#loc5309)
    %28369 = stablehlo.broadcast_in_dim %28368, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28370 = stablehlo.divide %25571, %28369 : tensor<1x2048xf32> loc(#loc5314)
    %28371 = stablehlo.multiply %28368, %28368 : tensor<f32> loc(#loc5310)
    %28372 = stablehlo.multiply %28370, %28370 : tensor<1x2048xf32> loc(#loc5310)
    %28373 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28374 = stablehlo.reduce(%28372 init: %28373) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28375 = stablehlo.multiply %28371, %28374 : tensor<f32> loc(#loc5310)
    %28376 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28377 = stablehlo.broadcast_in_dim %28376, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28378 = stablehlo.broadcast_in_dim %28375, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28379 = stablehlo.concatenate %28377, %28378, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28380 = stablehlo.abs %25572 : tensor<1x2048xf32> loc(#loc5307)
    %28381 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28382 = stablehlo.reduce(%28380 init: %28381) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28383 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28384 = stablehlo.maximum %28382, %28383 : tensor<f32> loc(#loc5309)
    %28385 = stablehlo.broadcast_in_dim %28384, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28386 = stablehlo.divide %25572, %28385 : tensor<1x2048xf32> loc(#loc5314)
    %28387 = stablehlo.multiply %28384, %28384 : tensor<f32> loc(#loc5310)
    %28388 = stablehlo.multiply %28386, %28386 : tensor<1x2048xf32> loc(#loc5310)
    %28389 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28390 = stablehlo.reduce(%28388 init: %28389) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28391 = stablehlo.multiply %28387, %28390 : tensor<f32> loc(#loc5310)
    %28392 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28393 = stablehlo.broadcast_in_dim %28392, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28394 = stablehlo.broadcast_in_dim %28391, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28395 = stablehlo.concatenate %28393, %28394, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28396 = stablehlo.abs %25573 : tensor<1x2048xf32> loc(#loc5307)
    %28397 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28398 = stablehlo.reduce(%28396 init: %28397) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28399 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28400 = stablehlo.maximum %28398, %28399 : tensor<f32> loc(#loc5309)
    %28401 = stablehlo.broadcast_in_dim %28400, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28402 = stablehlo.divide %25573, %28401 : tensor<1x2048xf32> loc(#loc5314)
    %28403 = stablehlo.multiply %28400, %28400 : tensor<f32> loc(#loc5310)
    %28404 = stablehlo.multiply %28402, %28402 : tensor<1x2048xf32> loc(#loc5310)
    %28405 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28406 = stablehlo.reduce(%28404 init: %28405) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28407 = stablehlo.multiply %28403, %28406 : tensor<f32> loc(#loc5310)
    %28408 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28409 = stablehlo.broadcast_in_dim %28408, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28410 = stablehlo.broadcast_in_dim %28407, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28411 = stablehlo.concatenate %28409, %28410, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28412 = stablehlo.abs %25574 : tensor<1x2048xf32> loc(#loc5307)
    %28413 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28414 = stablehlo.reduce(%28412 init: %28413) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28415 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28416 = stablehlo.maximum %28414, %28415 : tensor<f32> loc(#loc5309)
    %28417 = stablehlo.broadcast_in_dim %28416, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc5314)
    %28418 = stablehlo.divide %25574, %28417 : tensor<1x2048xf32> loc(#loc5314)
    %28419 = stablehlo.multiply %28416, %28416 : tensor<f32> loc(#loc5310)
    %28420 = stablehlo.multiply %28418, %28418 : tensor<1x2048xf32> loc(#loc5310)
    %28421 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5243)
    %28422 = stablehlo.reduce(%28420 init: %28421) across dimensions = [0, 1] : (tensor<1x2048xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5243)
      stablehlo.return %29929 : tensor<f32> loc(#loc5243)
    } loc(#loc5243)
    %28423 = stablehlo.multiply %28419, %28422 : tensor<f32> loc(#loc5310)
    %28424 = stablehlo.constant dense<2.048000e+03> : tensor<f32> loc(#loc)
    %28425 = stablehlo.broadcast_in_dim %28424, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28426 = stablehlo.broadcast_in_dim %28423, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28427 = stablehlo.concatenate %28425, %28426, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28428 = stablehlo.abs %25575 : tensor<1x3x2048x32x64xf32> loc(#loc5307)
    %28429 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5316)
    %28430 = stablehlo.reduce(%28428 init: %28429) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5316)
      stablehlo.return %29929 : tensor<f32> loc(#loc5316)
    } loc(#loc5316)
    %28431 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28432 = stablehlo.maximum %28430, %28431 : tensor<f32> loc(#loc5309)
    %28433 = stablehlo.broadcast_in_dim %28432, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %28434 = stablehlo.divide %25575, %28433 : tensor<1x3x2048x32x64xf32> loc(#loc5314)
    %28435 = stablehlo.multiply %28432, %28432 : tensor<f32> loc(#loc5310)
    %28436 = stablehlo.multiply %28434, %28434 : tensor<1x3x2048x32x64xf32> loc(#loc5310)
    %28437 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5244)
    %28438 = stablehlo.reduce(%28436 init: %28437) across dimensions = [0, 1, 2, 3, 4] : (tensor<1x3x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5244)
      stablehlo.return %29929 : tensor<f32> loc(#loc5244)
    } loc(#loc5244)
    %28439 = stablehlo.multiply %28435, %28438 : tensor<f32> loc(#loc5310)
    %28440 = stablehlo.constant dense<0x4B400000> : tensor<f32> loc(#loc)
    %28441 = stablehlo.broadcast_in_dim %28440, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28442 = stablehlo.broadcast_in_dim %28439, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28443 = stablehlo.concatenate %28441, %28442, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28444 = stablehlo.abs %25576 : tensor<1x64xf32> loc(#loc5307)
    %28445 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5313)
    %28446 = stablehlo.reduce(%28444 init: %28445) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5313)
      stablehlo.return %29929 : tensor<f32> loc(#loc5313)
    } loc(#loc5313)
    %28447 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28448 = stablehlo.maximum %28446, %28447 : tensor<f32> loc(#loc5309)
    %28449 = stablehlo.broadcast_in_dim %28448, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc5314)
    %28450 = stablehlo.divide %25576, %28449 : tensor<1x64xf32> loc(#loc5314)
    %28451 = stablehlo.multiply %28448, %28448 : tensor<f32> loc(#loc5310)
    %28452 = stablehlo.multiply %28450, %28450 : tensor<1x64xf32> loc(#loc5310)
    %28453 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5240)
    %28454 = stablehlo.reduce(%28452 init: %28453) across dimensions = [0, 1] : (tensor<1x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5240)
      stablehlo.return %29929 : tensor<f32> loc(#loc5240)
    } loc(#loc5240)
    %28455 = stablehlo.multiply %28451, %28454 : tensor<f32> loc(#loc5310)
    %28456 = stablehlo.constant dense<6.400000e+01> : tensor<f32> loc(#loc)
    %28457 = stablehlo.broadcast_in_dim %28456, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28458 = stablehlo.broadcast_in_dim %28455, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28459 = stablehlo.concatenate %28457, %28458, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28460 = stablehlo.abs %25577 : tensor<1x2048x32x64xf32> loc(#loc5307)
    %28461 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5317)
    %28462 = stablehlo.reduce(%28460 init: %28461) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.maximum %arg558, %arg559 : tensor<f32> loc(#loc5317)
      stablehlo.return %29929 : tensor<f32> loc(#loc5317)
    } loc(#loc5317)
    %28463 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %28464 = stablehlo.maximum %28462, %28463 : tensor<f32> loc(#loc5309)
    %28465 = stablehlo.broadcast_in_dim %28464, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc5314)
    %28466 = stablehlo.divide %25577, %28465 : tensor<1x2048x32x64xf32> loc(#loc5314)
    %28467 = stablehlo.multiply %28464, %28464 : tensor<f32> loc(#loc5310)
    %28468 = stablehlo.multiply %28466, %28466 : tensor<1x2048x32x64xf32> loc(#loc5310)
    %28469 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5245)
    %28470 = stablehlo.reduce(%28468 init: %28469) across dimensions = [0, 1, 2, 3] : (tensor<1x2048x32x64xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg558: tensor<f32> loc(unknown), %arg559: tensor<f32> loc(unknown))  {
      %29929 = stablehlo.add %arg558, %arg559 : tensor<f32> loc(#loc5245)
      stablehlo.return %29929 : tensor<f32> loc(#loc5245)
    } loc(#loc5245)
    %28471 = stablehlo.multiply %28467, %28470 : tensor<f32> loc(#loc5310)
    %28472 = stablehlo.constant dense<0x4A800000> : tensor<f32> loc(#loc)
    %28473 = stablehlo.broadcast_in_dim %28472, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28474 = stablehlo.broadcast_in_dim %28471, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5311)
    %28475 = stablehlo.concatenate %28473, %28474, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32> loc(#loc5312)
    %28476 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28477 = stablehlo.dynamic_slice %25595, %28476, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28478 = stablehlo.reshape %28477 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28479 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28480 = stablehlo.dynamic_slice %25595, %28479, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28481 = stablehlo.reshape %28480 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28482 = stablehlo.divide %28478, %28481 : tensor<f32> loc(#loc5253)
    %28483 = stablehlo.sqrt %28482 : tensor<f32> loc(#loc5249)
    %28484 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28485 = stablehlo.dynamic_slice %25611, %28484, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28486 = stablehlo.reshape %28485 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28487 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28488 = stablehlo.dynamic_slice %25611, %28487, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28489 = stablehlo.reshape %28488 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28490 = stablehlo.divide %28486, %28489 : tensor<f32> loc(#loc5253)
    %28491 = stablehlo.sqrt %28490 : tensor<f32> loc(#loc5249)
    %28492 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28493 = stablehlo.dynamic_slice %25627, %28492, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28494 = stablehlo.reshape %28493 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28495 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28496 = stablehlo.dynamic_slice %25627, %28495, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28497 = stablehlo.reshape %28496 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28498 = stablehlo.divide %28494, %28497 : tensor<f32> loc(#loc5253)
    %28499 = stablehlo.sqrt %28498 : tensor<f32> loc(#loc5249)
    %28500 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28501 = stablehlo.dynamic_slice %25643, %28500, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28502 = stablehlo.reshape %28501 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28503 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28504 = stablehlo.dynamic_slice %25643, %28503, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28505 = stablehlo.reshape %28504 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28506 = stablehlo.divide %28502, %28505 : tensor<f32> loc(#loc5253)
    %28507 = stablehlo.sqrt %28506 : tensor<f32> loc(#loc5249)
    %28508 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28509 = stablehlo.dynamic_slice %25659, %28508, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28510 = stablehlo.reshape %28509 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28511 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28512 = stablehlo.dynamic_slice %25659, %28511, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28513 = stablehlo.reshape %28512 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28514 = stablehlo.divide %28510, %28513 : tensor<f32> loc(#loc5253)
    %28515 = stablehlo.sqrt %28514 : tensor<f32> loc(#loc5249)
    %28516 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28517 = stablehlo.dynamic_slice %25675, %28516, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28518 = stablehlo.reshape %28517 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28519 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28520 = stablehlo.dynamic_slice %25675, %28519, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28521 = stablehlo.reshape %28520 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28522 = stablehlo.divide %28518, %28521 : tensor<f32> loc(#loc5253)
    %28523 = stablehlo.sqrt %28522 : tensor<f32> loc(#loc5249)
    %28524 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28525 = stablehlo.dynamic_slice %25691, %28524, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28526 = stablehlo.reshape %28525 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28527 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28528 = stablehlo.dynamic_slice %25691, %28527, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28529 = stablehlo.reshape %28528 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28530 = stablehlo.divide %28526, %28529 : tensor<f32> loc(#loc5253)
    %28531 = stablehlo.sqrt %28530 : tensor<f32> loc(#loc5249)
    %28532 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28533 = stablehlo.dynamic_slice %25707, %28532, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28534 = stablehlo.reshape %28533 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28535 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28536 = stablehlo.dynamic_slice %25707, %28535, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28537 = stablehlo.reshape %28536 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28538 = stablehlo.divide %28534, %28537 : tensor<f32> loc(#loc5253)
    %28539 = stablehlo.sqrt %28538 : tensor<f32> loc(#loc5249)
    %28540 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28541 = stablehlo.dynamic_slice %25723, %28540, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28542 = stablehlo.reshape %28541 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28543 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28544 = stablehlo.dynamic_slice %25723, %28543, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28545 = stablehlo.reshape %28544 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28546 = stablehlo.divide %28542, %28545 : tensor<f32> loc(#loc5253)
    %28547 = stablehlo.sqrt %28546 : tensor<f32> loc(#loc5249)
    %28548 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28549 = stablehlo.dynamic_slice %25739, %28548, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28550 = stablehlo.reshape %28549 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28551 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28552 = stablehlo.dynamic_slice %25739, %28551, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28553 = stablehlo.reshape %28552 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28554 = stablehlo.divide %28550, %28553 : tensor<f32> loc(#loc5253)
    %28555 = stablehlo.sqrt %28554 : tensor<f32> loc(#loc5249)
    %28556 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28557 = stablehlo.dynamic_slice %25755, %28556, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28558 = stablehlo.reshape %28557 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28559 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28560 = stablehlo.dynamic_slice %25755, %28559, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28561 = stablehlo.reshape %28560 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28562 = stablehlo.divide %28558, %28561 : tensor<f32> loc(#loc5253)
    %28563 = stablehlo.sqrt %28562 : tensor<f32> loc(#loc5249)
    %28564 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28565 = stablehlo.dynamic_slice %25771, %28564, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28566 = stablehlo.reshape %28565 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28567 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28568 = stablehlo.dynamic_slice %25771, %28567, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28569 = stablehlo.reshape %28568 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28570 = stablehlo.divide %28566, %28569 : tensor<f32> loc(#loc5253)
    %28571 = stablehlo.sqrt %28570 : tensor<f32> loc(#loc5249)
    %28572 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28573 = stablehlo.dynamic_slice %25787, %28572, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28574 = stablehlo.reshape %28573 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28575 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28576 = stablehlo.dynamic_slice %25787, %28575, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28577 = stablehlo.reshape %28576 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28578 = stablehlo.divide %28574, %28577 : tensor<f32> loc(#loc5253)
    %28579 = stablehlo.sqrt %28578 : tensor<f32> loc(#loc5249)
    %28580 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28581 = stablehlo.dynamic_slice %25803, %28580, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28582 = stablehlo.reshape %28581 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28583 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28584 = stablehlo.dynamic_slice %25803, %28583, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28585 = stablehlo.reshape %28584 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28586 = stablehlo.divide %28582, %28585 : tensor<f32> loc(#loc5253)
    %28587 = stablehlo.sqrt %28586 : tensor<f32> loc(#loc5249)
    %28588 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28589 = stablehlo.dynamic_slice %25819, %28588, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28590 = stablehlo.reshape %28589 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28591 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28592 = stablehlo.dynamic_slice %25819, %28591, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28593 = stablehlo.reshape %28592 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28594 = stablehlo.divide %28590, %28593 : tensor<f32> loc(#loc5253)
    %28595 = stablehlo.sqrt %28594 : tensor<f32> loc(#loc5249)
    %28596 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28597 = stablehlo.dynamic_slice %25835, %28596, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28598 = stablehlo.reshape %28597 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28599 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28600 = stablehlo.dynamic_slice %25835, %28599, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28601 = stablehlo.reshape %28600 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28602 = stablehlo.divide %28598, %28601 : tensor<f32> loc(#loc5253)
    %28603 = stablehlo.sqrt %28602 : tensor<f32> loc(#loc5249)
    %28604 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28605 = stablehlo.dynamic_slice %25851, %28604, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28606 = stablehlo.reshape %28605 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28607 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28608 = stablehlo.dynamic_slice %25851, %28607, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28609 = stablehlo.reshape %28608 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28610 = stablehlo.divide %28606, %28609 : tensor<f32> loc(#loc5253)
    %28611 = stablehlo.sqrt %28610 : tensor<f32> loc(#loc5249)
    %28612 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28613 = stablehlo.dynamic_slice %25867, %28612, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28614 = stablehlo.reshape %28613 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28615 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28616 = stablehlo.dynamic_slice %25867, %28615, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28617 = stablehlo.reshape %28616 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28618 = stablehlo.divide %28614, %28617 : tensor<f32> loc(#loc5253)
    %28619 = stablehlo.sqrt %28618 : tensor<f32> loc(#loc5249)
    %28620 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28621 = stablehlo.dynamic_slice %25883, %28620, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28622 = stablehlo.reshape %28621 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28623 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28624 = stablehlo.dynamic_slice %25883, %28623, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28625 = stablehlo.reshape %28624 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28626 = stablehlo.divide %28622, %28625 : tensor<f32> loc(#loc5253)
    %28627 = stablehlo.sqrt %28626 : tensor<f32> loc(#loc5249)
    %28628 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28629 = stablehlo.dynamic_slice %25899, %28628, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28630 = stablehlo.reshape %28629 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28631 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28632 = stablehlo.dynamic_slice %25899, %28631, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28633 = stablehlo.reshape %28632 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28634 = stablehlo.divide %28630, %28633 : tensor<f32> loc(#loc5253)
    %28635 = stablehlo.sqrt %28634 : tensor<f32> loc(#loc5249)
    %28636 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28637 = stablehlo.dynamic_slice %25915, %28636, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28638 = stablehlo.reshape %28637 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28639 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28640 = stablehlo.dynamic_slice %25915, %28639, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28641 = stablehlo.reshape %28640 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28642 = stablehlo.divide %28638, %28641 : tensor<f32> loc(#loc5253)
    %28643 = stablehlo.sqrt %28642 : tensor<f32> loc(#loc5249)
    %28644 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28645 = stablehlo.dynamic_slice %25931, %28644, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28646 = stablehlo.reshape %28645 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28647 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28648 = stablehlo.dynamic_slice %25931, %28647, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28649 = stablehlo.reshape %28648 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28650 = stablehlo.divide %28646, %28649 : tensor<f32> loc(#loc5253)
    %28651 = stablehlo.sqrt %28650 : tensor<f32> loc(#loc5249)
    %28652 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28653 = stablehlo.dynamic_slice %25947, %28652, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28654 = stablehlo.reshape %28653 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28655 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28656 = stablehlo.dynamic_slice %25947, %28655, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28657 = stablehlo.reshape %28656 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28658 = stablehlo.divide %28654, %28657 : tensor<f32> loc(#loc5253)
    %28659 = stablehlo.sqrt %28658 : tensor<f32> loc(#loc5249)
    %28660 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28661 = stablehlo.dynamic_slice %25963, %28660, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28662 = stablehlo.reshape %28661 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28663 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28664 = stablehlo.dynamic_slice %25963, %28663, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28665 = stablehlo.reshape %28664 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28666 = stablehlo.divide %28662, %28665 : tensor<f32> loc(#loc5253)
    %28667 = stablehlo.sqrt %28666 : tensor<f32> loc(#loc5249)
    %28668 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28669 = stablehlo.dynamic_slice %25979, %28668, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28670 = stablehlo.reshape %28669 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28671 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28672 = stablehlo.dynamic_slice %25979, %28671, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28673 = stablehlo.reshape %28672 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28674 = stablehlo.divide %28670, %28673 : tensor<f32> loc(#loc5253)
    %28675 = stablehlo.sqrt %28674 : tensor<f32> loc(#loc5249)
    %28676 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28677 = stablehlo.dynamic_slice %25995, %28676, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28678 = stablehlo.reshape %28677 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28679 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28680 = stablehlo.dynamic_slice %25995, %28679, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28681 = stablehlo.reshape %28680 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28682 = stablehlo.divide %28678, %28681 : tensor<f32> loc(#loc5253)
    %28683 = stablehlo.sqrt %28682 : tensor<f32> loc(#loc5249)
    %28684 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28685 = stablehlo.dynamic_slice %26011, %28684, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28686 = stablehlo.reshape %28685 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28687 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28688 = stablehlo.dynamic_slice %26011, %28687, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28689 = stablehlo.reshape %28688 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28690 = stablehlo.divide %28686, %28689 : tensor<f32> loc(#loc5253)
    %28691 = stablehlo.sqrt %28690 : tensor<f32> loc(#loc5249)
    %28692 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28693 = stablehlo.dynamic_slice %26027, %28692, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28694 = stablehlo.reshape %28693 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28695 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28696 = stablehlo.dynamic_slice %26027, %28695, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28697 = stablehlo.reshape %28696 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28698 = stablehlo.divide %28694, %28697 : tensor<f32> loc(#loc5253)
    %28699 = stablehlo.sqrt %28698 : tensor<f32> loc(#loc5249)
    %28700 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28701 = stablehlo.dynamic_slice %26043, %28700, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28702 = stablehlo.reshape %28701 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28703 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28704 = stablehlo.dynamic_slice %26043, %28703, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28705 = stablehlo.reshape %28704 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28706 = stablehlo.divide %28702, %28705 : tensor<f32> loc(#loc5253)
    %28707 = stablehlo.sqrt %28706 : tensor<f32> loc(#loc5249)
    %28708 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28709 = stablehlo.dynamic_slice %26059, %28708, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28710 = stablehlo.reshape %28709 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28711 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28712 = stablehlo.dynamic_slice %26059, %28711, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28713 = stablehlo.reshape %28712 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28714 = stablehlo.divide %28710, %28713 : tensor<f32> loc(#loc5253)
    %28715 = stablehlo.sqrt %28714 : tensor<f32> loc(#loc5249)
    %28716 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28717 = stablehlo.dynamic_slice %26075, %28716, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28718 = stablehlo.reshape %28717 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28719 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28720 = stablehlo.dynamic_slice %26075, %28719, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28721 = stablehlo.reshape %28720 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28722 = stablehlo.divide %28718, %28721 : tensor<f32> loc(#loc5253)
    %28723 = stablehlo.sqrt %28722 : tensor<f32> loc(#loc5249)
    %28724 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28725 = stablehlo.dynamic_slice %26091, %28724, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28726 = stablehlo.reshape %28725 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28727 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28728 = stablehlo.dynamic_slice %26091, %28727, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28729 = stablehlo.reshape %28728 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28730 = stablehlo.divide %28726, %28729 : tensor<f32> loc(#loc5253)
    %28731 = stablehlo.sqrt %28730 : tensor<f32> loc(#loc5249)
    %28732 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28733 = stablehlo.dynamic_slice %26107, %28732, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28734 = stablehlo.reshape %28733 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28735 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28736 = stablehlo.dynamic_slice %26107, %28735, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28737 = stablehlo.reshape %28736 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28738 = stablehlo.divide %28734, %28737 : tensor<f32> loc(#loc5253)
    %28739 = stablehlo.sqrt %28738 : tensor<f32> loc(#loc5249)
    %28740 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28741 = stablehlo.dynamic_slice %26123, %28740, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28742 = stablehlo.reshape %28741 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28743 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28744 = stablehlo.dynamic_slice %26123, %28743, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28745 = stablehlo.reshape %28744 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28746 = stablehlo.divide %28742, %28745 : tensor<f32> loc(#loc5253)
    %28747 = stablehlo.sqrt %28746 : tensor<f32> loc(#loc5249)
    %28748 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28749 = stablehlo.dynamic_slice %26139, %28748, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28750 = stablehlo.reshape %28749 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28751 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28752 = stablehlo.dynamic_slice %26139, %28751, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28753 = stablehlo.reshape %28752 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28754 = stablehlo.divide %28750, %28753 : tensor<f32> loc(#loc5253)
    %28755 = stablehlo.sqrt %28754 : tensor<f32> loc(#loc5249)
    %28756 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28757 = stablehlo.dynamic_slice %26155, %28756, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28758 = stablehlo.reshape %28757 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28759 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28760 = stablehlo.dynamic_slice %26155, %28759, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28761 = stablehlo.reshape %28760 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28762 = stablehlo.divide %28758, %28761 : tensor<f32> loc(#loc5253)
    %28763 = stablehlo.sqrt %28762 : tensor<f32> loc(#loc5249)
    %28764 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28765 = stablehlo.dynamic_slice %26171, %28764, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28766 = stablehlo.reshape %28765 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28767 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28768 = stablehlo.dynamic_slice %26171, %28767, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28769 = stablehlo.reshape %28768 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28770 = stablehlo.divide %28766, %28769 : tensor<f32> loc(#loc5253)
    %28771 = stablehlo.sqrt %28770 : tensor<f32> loc(#loc5249)
    %28772 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28773 = stablehlo.dynamic_slice %26187, %28772, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28774 = stablehlo.reshape %28773 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28775 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28776 = stablehlo.dynamic_slice %26187, %28775, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28777 = stablehlo.reshape %28776 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28778 = stablehlo.divide %28774, %28777 : tensor<f32> loc(#loc5253)
    %28779 = stablehlo.sqrt %28778 : tensor<f32> loc(#loc5249)
    %28780 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28781 = stablehlo.dynamic_slice %26203, %28780, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28782 = stablehlo.reshape %28781 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28783 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28784 = stablehlo.dynamic_slice %26203, %28783, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28785 = stablehlo.reshape %28784 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28786 = stablehlo.divide %28782, %28785 : tensor<f32> loc(#loc5253)
    %28787 = stablehlo.sqrt %28786 : tensor<f32> loc(#loc5249)
    %28788 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28789 = stablehlo.dynamic_slice %26219, %28788, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28790 = stablehlo.reshape %28789 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28791 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28792 = stablehlo.dynamic_slice %26219, %28791, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28793 = stablehlo.reshape %28792 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28794 = stablehlo.divide %28790, %28793 : tensor<f32> loc(#loc5253)
    %28795 = stablehlo.sqrt %28794 : tensor<f32> loc(#loc5249)
    %28796 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28797 = stablehlo.dynamic_slice %26235, %28796, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28798 = stablehlo.reshape %28797 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28799 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28800 = stablehlo.dynamic_slice %26235, %28799, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28801 = stablehlo.reshape %28800 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28802 = stablehlo.divide %28798, %28801 : tensor<f32> loc(#loc5253)
    %28803 = stablehlo.sqrt %28802 : tensor<f32> loc(#loc5249)
    %28804 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28805 = stablehlo.dynamic_slice %26251, %28804, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28806 = stablehlo.reshape %28805 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28807 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28808 = stablehlo.dynamic_slice %26251, %28807, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28809 = stablehlo.reshape %28808 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28810 = stablehlo.divide %28806, %28809 : tensor<f32> loc(#loc5253)
    %28811 = stablehlo.sqrt %28810 : tensor<f32> loc(#loc5249)
    %28812 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28813 = stablehlo.dynamic_slice %26267, %28812, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28814 = stablehlo.reshape %28813 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28815 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28816 = stablehlo.dynamic_slice %26267, %28815, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28817 = stablehlo.reshape %28816 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28818 = stablehlo.divide %28814, %28817 : tensor<f32> loc(#loc5253)
    %28819 = stablehlo.sqrt %28818 : tensor<f32> loc(#loc5249)
    %28820 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28821 = stablehlo.dynamic_slice %26283, %28820, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28822 = stablehlo.reshape %28821 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28823 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28824 = stablehlo.dynamic_slice %26283, %28823, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28825 = stablehlo.reshape %28824 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28826 = stablehlo.divide %28822, %28825 : tensor<f32> loc(#loc5253)
    %28827 = stablehlo.sqrt %28826 : tensor<f32> loc(#loc5249)
    %28828 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28829 = stablehlo.dynamic_slice %26299, %28828, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28830 = stablehlo.reshape %28829 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28831 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28832 = stablehlo.dynamic_slice %26299, %28831, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28833 = stablehlo.reshape %28832 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28834 = stablehlo.divide %28830, %28833 : tensor<f32> loc(#loc5253)
    %28835 = stablehlo.sqrt %28834 : tensor<f32> loc(#loc5249)
    %28836 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28837 = stablehlo.dynamic_slice %26315, %28836, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28838 = stablehlo.reshape %28837 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28839 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28840 = stablehlo.dynamic_slice %26315, %28839, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28841 = stablehlo.reshape %28840 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28842 = stablehlo.divide %28838, %28841 : tensor<f32> loc(#loc5253)
    %28843 = stablehlo.sqrt %28842 : tensor<f32> loc(#loc5249)
    %28844 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28845 = stablehlo.dynamic_slice %26331, %28844, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28846 = stablehlo.reshape %28845 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28847 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28848 = stablehlo.dynamic_slice %26331, %28847, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28849 = stablehlo.reshape %28848 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28850 = stablehlo.divide %28846, %28849 : tensor<f32> loc(#loc5253)
    %28851 = stablehlo.sqrt %28850 : tensor<f32> loc(#loc5249)
    %28852 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28853 = stablehlo.dynamic_slice %26347, %28852, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28854 = stablehlo.reshape %28853 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28855 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28856 = stablehlo.dynamic_slice %26347, %28855, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28857 = stablehlo.reshape %28856 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28858 = stablehlo.divide %28854, %28857 : tensor<f32> loc(#loc5253)
    %28859 = stablehlo.sqrt %28858 : tensor<f32> loc(#loc5249)
    %28860 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28861 = stablehlo.dynamic_slice %26363, %28860, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28862 = stablehlo.reshape %28861 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28863 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28864 = stablehlo.dynamic_slice %26363, %28863, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28865 = stablehlo.reshape %28864 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28866 = stablehlo.divide %28862, %28865 : tensor<f32> loc(#loc5253)
    %28867 = stablehlo.sqrt %28866 : tensor<f32> loc(#loc5249)
    %28868 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28869 = stablehlo.dynamic_slice %26379, %28868, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28870 = stablehlo.reshape %28869 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28871 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28872 = stablehlo.dynamic_slice %26379, %28871, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28873 = stablehlo.reshape %28872 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28874 = stablehlo.divide %28870, %28873 : tensor<f32> loc(#loc5253)
    %28875 = stablehlo.sqrt %28874 : tensor<f32> loc(#loc5249)
    %28876 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28877 = stablehlo.dynamic_slice %26395, %28876, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28878 = stablehlo.reshape %28877 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28879 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28880 = stablehlo.dynamic_slice %26395, %28879, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28881 = stablehlo.reshape %28880 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28882 = stablehlo.divide %28878, %28881 : tensor<f32> loc(#loc5253)
    %28883 = stablehlo.sqrt %28882 : tensor<f32> loc(#loc5249)
    %28884 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28885 = stablehlo.dynamic_slice %26411, %28884, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28886 = stablehlo.reshape %28885 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28887 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28888 = stablehlo.dynamic_slice %26411, %28887, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28889 = stablehlo.reshape %28888 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28890 = stablehlo.divide %28886, %28889 : tensor<f32> loc(#loc5253)
    %28891 = stablehlo.sqrt %28890 : tensor<f32> loc(#loc5249)
    %28892 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28893 = stablehlo.dynamic_slice %26427, %28892, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28894 = stablehlo.reshape %28893 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28895 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28896 = stablehlo.dynamic_slice %26427, %28895, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28897 = stablehlo.reshape %28896 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28898 = stablehlo.divide %28894, %28897 : tensor<f32> loc(#loc5253)
    %28899 = stablehlo.sqrt %28898 : tensor<f32> loc(#loc5249)
    %28900 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28901 = stablehlo.dynamic_slice %26443, %28900, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28902 = stablehlo.reshape %28901 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28903 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28904 = stablehlo.dynamic_slice %26443, %28903, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28905 = stablehlo.reshape %28904 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28906 = stablehlo.divide %28902, %28905 : tensor<f32> loc(#loc5253)
    %28907 = stablehlo.sqrt %28906 : tensor<f32> loc(#loc5249)
    %28908 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28909 = stablehlo.dynamic_slice %26459, %28908, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28910 = stablehlo.reshape %28909 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28911 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28912 = stablehlo.dynamic_slice %26459, %28911, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28913 = stablehlo.reshape %28912 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28914 = stablehlo.divide %28910, %28913 : tensor<f32> loc(#loc5253)
    %28915 = stablehlo.sqrt %28914 : tensor<f32> loc(#loc5249)
    %28916 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28917 = stablehlo.dynamic_slice %26475, %28916, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28918 = stablehlo.reshape %28917 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28919 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28920 = stablehlo.dynamic_slice %26475, %28919, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28921 = stablehlo.reshape %28920 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28922 = stablehlo.divide %28918, %28921 : tensor<f32> loc(#loc5253)
    %28923 = stablehlo.sqrt %28922 : tensor<f32> loc(#loc5249)
    %28924 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28925 = stablehlo.dynamic_slice %26491, %28924, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28926 = stablehlo.reshape %28925 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28927 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28928 = stablehlo.dynamic_slice %26491, %28927, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28929 = stablehlo.reshape %28928 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28930 = stablehlo.divide %28926, %28929 : tensor<f32> loc(#loc5253)
    %28931 = stablehlo.sqrt %28930 : tensor<f32> loc(#loc5249)
    %28932 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28933 = stablehlo.dynamic_slice %26507, %28932, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28934 = stablehlo.reshape %28933 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28935 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28936 = stablehlo.dynamic_slice %26507, %28935, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28937 = stablehlo.reshape %28936 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28938 = stablehlo.divide %28934, %28937 : tensor<f32> loc(#loc5253)
    %28939 = stablehlo.sqrt %28938 : tensor<f32> loc(#loc5249)
    %28940 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28941 = stablehlo.dynamic_slice %26523, %28940, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28942 = stablehlo.reshape %28941 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28943 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28944 = stablehlo.dynamic_slice %26523, %28943, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28945 = stablehlo.reshape %28944 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28946 = stablehlo.divide %28942, %28945 : tensor<f32> loc(#loc5253)
    %28947 = stablehlo.sqrt %28946 : tensor<f32> loc(#loc5249)
    %28948 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28949 = stablehlo.dynamic_slice %26539, %28948, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28950 = stablehlo.reshape %28949 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28951 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28952 = stablehlo.dynamic_slice %26539, %28951, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28953 = stablehlo.reshape %28952 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28954 = stablehlo.divide %28950, %28953 : tensor<f32> loc(#loc5253)
    %28955 = stablehlo.sqrt %28954 : tensor<f32> loc(#loc5249)
    %28956 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28957 = stablehlo.dynamic_slice %26555, %28956, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28958 = stablehlo.reshape %28957 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28959 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28960 = stablehlo.dynamic_slice %26555, %28959, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28961 = stablehlo.reshape %28960 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28962 = stablehlo.divide %28958, %28961 : tensor<f32> loc(#loc5253)
    %28963 = stablehlo.sqrt %28962 : tensor<f32> loc(#loc5249)
    %28964 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28965 = stablehlo.dynamic_slice %26571, %28964, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28966 = stablehlo.reshape %28965 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28967 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28968 = stablehlo.dynamic_slice %26571, %28967, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28969 = stablehlo.reshape %28968 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28970 = stablehlo.divide %28966, %28969 : tensor<f32> loc(#loc5253)
    %28971 = stablehlo.sqrt %28970 : tensor<f32> loc(#loc5249)
    %28972 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28973 = stablehlo.dynamic_slice %26587, %28972, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28974 = stablehlo.reshape %28973 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28975 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28976 = stablehlo.dynamic_slice %26587, %28975, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28977 = stablehlo.reshape %28976 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28978 = stablehlo.divide %28974, %28977 : tensor<f32> loc(#loc5253)
    %28979 = stablehlo.sqrt %28978 : tensor<f32> loc(#loc5249)
    %28980 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28981 = stablehlo.dynamic_slice %26603, %28980, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28982 = stablehlo.reshape %28981 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28983 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28984 = stablehlo.dynamic_slice %26603, %28983, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28985 = stablehlo.reshape %28984 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28986 = stablehlo.divide %28982, %28985 : tensor<f32> loc(#loc5253)
    %28987 = stablehlo.sqrt %28986 : tensor<f32> loc(#loc5249)
    %28988 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28989 = stablehlo.dynamic_slice %26619, %28988, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28990 = stablehlo.reshape %28989 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28991 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %28992 = stablehlo.dynamic_slice %26619, %28991, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28993 = stablehlo.reshape %28992 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28994 = stablehlo.divide %28990, %28993 : tensor<f32> loc(#loc5253)
    %28995 = stablehlo.sqrt %28994 : tensor<f32> loc(#loc5249)
    %28996 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %28997 = stablehlo.dynamic_slice %26635, %28996, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %28998 = stablehlo.reshape %28997 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %28999 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29000 = stablehlo.dynamic_slice %26635, %28999, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29001 = stablehlo.reshape %29000 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29002 = stablehlo.divide %28998, %29001 : tensor<f32> loc(#loc5253)
    %29003 = stablehlo.sqrt %29002 : tensor<f32> loc(#loc5249)
    %29004 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29005 = stablehlo.dynamic_slice %26651, %29004, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29006 = stablehlo.reshape %29005 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29007 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29008 = stablehlo.dynamic_slice %26651, %29007, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29009 = stablehlo.reshape %29008 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29010 = stablehlo.divide %29006, %29009 : tensor<f32> loc(#loc5253)
    %29011 = stablehlo.sqrt %29010 : tensor<f32> loc(#loc5249)
    %29012 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29013 = stablehlo.dynamic_slice %26667, %29012, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29014 = stablehlo.reshape %29013 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29015 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29016 = stablehlo.dynamic_slice %26667, %29015, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29017 = stablehlo.reshape %29016 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29018 = stablehlo.divide %29014, %29017 : tensor<f32> loc(#loc5253)
    %29019 = stablehlo.sqrt %29018 : tensor<f32> loc(#loc5249)
    %29020 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29021 = stablehlo.dynamic_slice %26683, %29020, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29022 = stablehlo.reshape %29021 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29023 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29024 = stablehlo.dynamic_slice %26683, %29023, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29025 = stablehlo.reshape %29024 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29026 = stablehlo.divide %29022, %29025 : tensor<f32> loc(#loc5253)
    %29027 = stablehlo.sqrt %29026 : tensor<f32> loc(#loc5249)
    %29028 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29029 = stablehlo.dynamic_slice %26699, %29028, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29030 = stablehlo.reshape %29029 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29031 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29032 = stablehlo.dynamic_slice %26699, %29031, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29033 = stablehlo.reshape %29032 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29034 = stablehlo.divide %29030, %29033 : tensor<f32> loc(#loc5253)
    %29035 = stablehlo.sqrt %29034 : tensor<f32> loc(#loc5249)
    %29036 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29037 = stablehlo.dynamic_slice %26715, %29036, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29038 = stablehlo.reshape %29037 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29039 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29040 = stablehlo.dynamic_slice %26715, %29039, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29041 = stablehlo.reshape %29040 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29042 = stablehlo.divide %29038, %29041 : tensor<f32> loc(#loc5253)
    %29043 = stablehlo.sqrt %29042 : tensor<f32> loc(#loc5249)
    %29044 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29045 = stablehlo.dynamic_slice %26731, %29044, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29046 = stablehlo.reshape %29045 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29047 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29048 = stablehlo.dynamic_slice %26731, %29047, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29049 = stablehlo.reshape %29048 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29050 = stablehlo.divide %29046, %29049 : tensor<f32> loc(#loc5253)
    %29051 = stablehlo.sqrt %29050 : tensor<f32> loc(#loc5249)
    %29052 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29053 = stablehlo.dynamic_slice %26747, %29052, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29054 = stablehlo.reshape %29053 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29055 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29056 = stablehlo.dynamic_slice %26747, %29055, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29057 = stablehlo.reshape %29056 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29058 = stablehlo.divide %29054, %29057 : tensor<f32> loc(#loc5253)
    %29059 = stablehlo.sqrt %29058 : tensor<f32> loc(#loc5249)
    %29060 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29061 = stablehlo.dynamic_slice %26763, %29060, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29062 = stablehlo.reshape %29061 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29063 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29064 = stablehlo.dynamic_slice %26763, %29063, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29065 = stablehlo.reshape %29064 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29066 = stablehlo.divide %29062, %29065 : tensor<f32> loc(#loc5253)
    %29067 = stablehlo.sqrt %29066 : tensor<f32> loc(#loc5249)
    %29068 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29069 = stablehlo.dynamic_slice %26779, %29068, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29070 = stablehlo.reshape %29069 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29071 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29072 = stablehlo.dynamic_slice %26779, %29071, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29073 = stablehlo.reshape %29072 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29074 = stablehlo.divide %29070, %29073 : tensor<f32> loc(#loc5253)
    %29075 = stablehlo.sqrt %29074 : tensor<f32> loc(#loc5249)
    %29076 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29077 = stablehlo.dynamic_slice %26795, %29076, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29078 = stablehlo.reshape %29077 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29079 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29080 = stablehlo.dynamic_slice %26795, %29079, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29081 = stablehlo.reshape %29080 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29082 = stablehlo.divide %29078, %29081 : tensor<f32> loc(#loc5253)
    %29083 = stablehlo.sqrt %29082 : tensor<f32> loc(#loc5249)
    %29084 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29085 = stablehlo.dynamic_slice %26811, %29084, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29086 = stablehlo.reshape %29085 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29087 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29088 = stablehlo.dynamic_slice %26811, %29087, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29089 = stablehlo.reshape %29088 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29090 = stablehlo.divide %29086, %29089 : tensor<f32> loc(#loc5253)
    %29091 = stablehlo.sqrt %29090 : tensor<f32> loc(#loc5249)
    %29092 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29093 = stablehlo.dynamic_slice %26827, %29092, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29094 = stablehlo.reshape %29093 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29095 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29096 = stablehlo.dynamic_slice %26827, %29095, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29097 = stablehlo.reshape %29096 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29098 = stablehlo.divide %29094, %29097 : tensor<f32> loc(#loc5253)
    %29099 = stablehlo.sqrt %29098 : tensor<f32> loc(#loc5249)
    %29100 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29101 = stablehlo.dynamic_slice %26843, %29100, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29102 = stablehlo.reshape %29101 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29103 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29104 = stablehlo.dynamic_slice %26843, %29103, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29105 = stablehlo.reshape %29104 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29106 = stablehlo.divide %29102, %29105 : tensor<f32> loc(#loc5253)
    %29107 = stablehlo.sqrt %29106 : tensor<f32> loc(#loc5249)
    %29108 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29109 = stablehlo.dynamic_slice %26859, %29108, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29110 = stablehlo.reshape %29109 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29111 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29112 = stablehlo.dynamic_slice %26859, %29111, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29113 = stablehlo.reshape %29112 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29114 = stablehlo.divide %29110, %29113 : tensor<f32> loc(#loc5253)
    %29115 = stablehlo.sqrt %29114 : tensor<f32> loc(#loc5249)
    %29116 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29117 = stablehlo.dynamic_slice %26875, %29116, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29118 = stablehlo.reshape %29117 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29119 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29120 = stablehlo.dynamic_slice %26875, %29119, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29121 = stablehlo.reshape %29120 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29122 = stablehlo.divide %29118, %29121 : tensor<f32> loc(#loc5253)
    %29123 = stablehlo.sqrt %29122 : tensor<f32> loc(#loc5249)
    %29124 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29125 = stablehlo.dynamic_slice %26891, %29124, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29126 = stablehlo.reshape %29125 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29127 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29128 = stablehlo.dynamic_slice %26891, %29127, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29129 = stablehlo.reshape %29128 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29130 = stablehlo.divide %29126, %29129 : tensor<f32> loc(#loc5253)
    %29131 = stablehlo.sqrt %29130 : tensor<f32> loc(#loc5249)
    %29132 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29133 = stablehlo.dynamic_slice %26907, %29132, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29134 = stablehlo.reshape %29133 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29135 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29136 = stablehlo.dynamic_slice %26907, %29135, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29137 = stablehlo.reshape %29136 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29138 = stablehlo.divide %29134, %29137 : tensor<f32> loc(#loc5253)
    %29139 = stablehlo.sqrt %29138 : tensor<f32> loc(#loc5249)
    %29140 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29141 = stablehlo.dynamic_slice %26923, %29140, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29142 = stablehlo.reshape %29141 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29143 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29144 = stablehlo.dynamic_slice %26923, %29143, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29145 = stablehlo.reshape %29144 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29146 = stablehlo.divide %29142, %29145 : tensor<f32> loc(#loc5253)
    %29147 = stablehlo.sqrt %29146 : tensor<f32> loc(#loc5249)
    %29148 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29149 = stablehlo.dynamic_slice %26939, %29148, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29150 = stablehlo.reshape %29149 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29151 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29152 = stablehlo.dynamic_slice %26939, %29151, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29153 = stablehlo.reshape %29152 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29154 = stablehlo.divide %29150, %29153 : tensor<f32> loc(#loc5253)
    %29155 = stablehlo.sqrt %29154 : tensor<f32> loc(#loc5249)
    %29156 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29157 = stablehlo.dynamic_slice %26955, %29156, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29158 = stablehlo.reshape %29157 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29159 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29160 = stablehlo.dynamic_slice %26955, %29159, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29161 = stablehlo.reshape %29160 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29162 = stablehlo.divide %29158, %29161 : tensor<f32> loc(#loc5253)
    %29163 = stablehlo.sqrt %29162 : tensor<f32> loc(#loc5249)
    %29164 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29165 = stablehlo.dynamic_slice %26971, %29164, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29166 = stablehlo.reshape %29165 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29167 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29168 = stablehlo.dynamic_slice %26971, %29167, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29169 = stablehlo.reshape %29168 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29170 = stablehlo.divide %29166, %29169 : tensor<f32> loc(#loc5253)
    %29171 = stablehlo.sqrt %29170 : tensor<f32> loc(#loc5249)
    %29172 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29173 = stablehlo.dynamic_slice %26987, %29172, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29174 = stablehlo.reshape %29173 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29175 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29176 = stablehlo.dynamic_slice %26987, %29175, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29177 = stablehlo.reshape %29176 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29178 = stablehlo.divide %29174, %29177 : tensor<f32> loc(#loc5253)
    %29179 = stablehlo.sqrt %29178 : tensor<f32> loc(#loc5249)
    %29180 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29181 = stablehlo.dynamic_slice %27003, %29180, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29182 = stablehlo.reshape %29181 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29183 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29184 = stablehlo.dynamic_slice %27003, %29183, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29185 = stablehlo.reshape %29184 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29186 = stablehlo.divide %29182, %29185 : tensor<f32> loc(#loc5253)
    %29187 = stablehlo.sqrt %29186 : tensor<f32> loc(#loc5249)
    %29188 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29189 = stablehlo.dynamic_slice %27019, %29188, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29190 = stablehlo.reshape %29189 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29191 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29192 = stablehlo.dynamic_slice %27019, %29191, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29193 = stablehlo.reshape %29192 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29194 = stablehlo.divide %29190, %29193 : tensor<f32> loc(#loc5253)
    %29195 = stablehlo.sqrt %29194 : tensor<f32> loc(#loc5249)
    %29196 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29197 = stablehlo.dynamic_slice %27035, %29196, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29198 = stablehlo.reshape %29197 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29199 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29200 = stablehlo.dynamic_slice %27035, %29199, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29201 = stablehlo.reshape %29200 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29202 = stablehlo.divide %29198, %29201 : tensor<f32> loc(#loc5253)
    %29203 = stablehlo.sqrt %29202 : tensor<f32> loc(#loc5249)
    %29204 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29205 = stablehlo.dynamic_slice %27051, %29204, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29206 = stablehlo.reshape %29205 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29207 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29208 = stablehlo.dynamic_slice %27051, %29207, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29209 = stablehlo.reshape %29208 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29210 = stablehlo.divide %29206, %29209 : tensor<f32> loc(#loc5253)
    %29211 = stablehlo.sqrt %29210 : tensor<f32> loc(#loc5249)
    %29212 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29213 = stablehlo.dynamic_slice %27067, %29212, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29214 = stablehlo.reshape %29213 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29215 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29216 = stablehlo.dynamic_slice %27067, %29215, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29217 = stablehlo.reshape %29216 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29218 = stablehlo.divide %29214, %29217 : tensor<f32> loc(#loc5253)
    %29219 = stablehlo.sqrt %29218 : tensor<f32> loc(#loc5249)
    %29220 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29221 = stablehlo.dynamic_slice %27083, %29220, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29222 = stablehlo.reshape %29221 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29223 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29224 = stablehlo.dynamic_slice %27083, %29223, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29225 = stablehlo.reshape %29224 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29226 = stablehlo.divide %29222, %29225 : tensor<f32> loc(#loc5253)
    %29227 = stablehlo.sqrt %29226 : tensor<f32> loc(#loc5249)
    %29228 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29229 = stablehlo.dynamic_slice %27099, %29228, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29230 = stablehlo.reshape %29229 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29231 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29232 = stablehlo.dynamic_slice %27099, %29231, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29233 = stablehlo.reshape %29232 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29234 = stablehlo.divide %29230, %29233 : tensor<f32> loc(#loc5253)
    %29235 = stablehlo.sqrt %29234 : tensor<f32> loc(#loc5249)
    %29236 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29237 = stablehlo.dynamic_slice %27115, %29236, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29238 = stablehlo.reshape %29237 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29239 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29240 = stablehlo.dynamic_slice %27115, %29239, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29241 = stablehlo.reshape %29240 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29242 = stablehlo.divide %29238, %29241 : tensor<f32> loc(#loc5253)
    %29243 = stablehlo.sqrt %29242 : tensor<f32> loc(#loc5249)
    %29244 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29245 = stablehlo.dynamic_slice %27131, %29244, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29246 = stablehlo.reshape %29245 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29247 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29248 = stablehlo.dynamic_slice %27131, %29247, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29249 = stablehlo.reshape %29248 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29250 = stablehlo.divide %29246, %29249 : tensor<f32> loc(#loc5253)
    %29251 = stablehlo.sqrt %29250 : tensor<f32> loc(#loc5249)
    %29252 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29253 = stablehlo.dynamic_slice %27147, %29252, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29254 = stablehlo.reshape %29253 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29255 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29256 = stablehlo.dynamic_slice %27147, %29255, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29257 = stablehlo.reshape %29256 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29258 = stablehlo.divide %29254, %29257 : tensor<f32> loc(#loc5253)
    %29259 = stablehlo.sqrt %29258 : tensor<f32> loc(#loc5249)
    %29260 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29261 = stablehlo.dynamic_slice %27163, %29260, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29262 = stablehlo.reshape %29261 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29263 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29264 = stablehlo.dynamic_slice %27163, %29263, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29265 = stablehlo.reshape %29264 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29266 = stablehlo.divide %29262, %29265 : tensor<f32> loc(#loc5253)
    %29267 = stablehlo.sqrt %29266 : tensor<f32> loc(#loc5249)
    %29268 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29269 = stablehlo.dynamic_slice %27179, %29268, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29270 = stablehlo.reshape %29269 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29271 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29272 = stablehlo.dynamic_slice %27179, %29271, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29273 = stablehlo.reshape %29272 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29274 = stablehlo.divide %29270, %29273 : tensor<f32> loc(#loc5253)
    %29275 = stablehlo.sqrt %29274 : tensor<f32> loc(#loc5249)
    %29276 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29277 = stablehlo.dynamic_slice %27195, %29276, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29278 = stablehlo.reshape %29277 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29279 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29280 = stablehlo.dynamic_slice %27195, %29279, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29281 = stablehlo.reshape %29280 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29282 = stablehlo.divide %29278, %29281 : tensor<f32> loc(#loc5253)
    %29283 = stablehlo.sqrt %29282 : tensor<f32> loc(#loc5249)
    %29284 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29285 = stablehlo.dynamic_slice %27211, %29284, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29286 = stablehlo.reshape %29285 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29287 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29288 = stablehlo.dynamic_slice %27211, %29287, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29289 = stablehlo.reshape %29288 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29290 = stablehlo.divide %29286, %29289 : tensor<f32> loc(#loc5253)
    %29291 = stablehlo.sqrt %29290 : tensor<f32> loc(#loc5249)
    %29292 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29293 = stablehlo.dynamic_slice %27227, %29292, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29294 = stablehlo.reshape %29293 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29295 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29296 = stablehlo.dynamic_slice %27227, %29295, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29297 = stablehlo.reshape %29296 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29298 = stablehlo.divide %29294, %29297 : tensor<f32> loc(#loc5253)
    %29299 = stablehlo.sqrt %29298 : tensor<f32> loc(#loc5249)
    %29300 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29301 = stablehlo.dynamic_slice %27243, %29300, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29302 = stablehlo.reshape %29301 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29303 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29304 = stablehlo.dynamic_slice %27243, %29303, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29305 = stablehlo.reshape %29304 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29306 = stablehlo.divide %29302, %29305 : tensor<f32> loc(#loc5253)
    %29307 = stablehlo.sqrt %29306 : tensor<f32> loc(#loc5249)
    %29308 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29309 = stablehlo.dynamic_slice %27259, %29308, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29310 = stablehlo.reshape %29309 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29311 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29312 = stablehlo.dynamic_slice %27259, %29311, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29313 = stablehlo.reshape %29312 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29314 = stablehlo.divide %29310, %29313 : tensor<f32> loc(#loc5253)
    %29315 = stablehlo.sqrt %29314 : tensor<f32> loc(#loc5249)
    %29316 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29317 = stablehlo.dynamic_slice %27275, %29316, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29318 = stablehlo.reshape %29317 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29319 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29320 = stablehlo.dynamic_slice %27275, %29319, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29321 = stablehlo.reshape %29320 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29322 = stablehlo.divide %29318, %29321 : tensor<f32> loc(#loc5253)
    %29323 = stablehlo.sqrt %29322 : tensor<f32> loc(#loc5249)
    %29324 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29325 = stablehlo.dynamic_slice %27291, %29324, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29326 = stablehlo.reshape %29325 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29327 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29328 = stablehlo.dynamic_slice %27291, %29327, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29329 = stablehlo.reshape %29328 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29330 = stablehlo.divide %29326, %29329 : tensor<f32> loc(#loc5253)
    %29331 = stablehlo.sqrt %29330 : tensor<f32> loc(#loc5249)
    %29332 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29333 = stablehlo.dynamic_slice %27307, %29332, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29334 = stablehlo.reshape %29333 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29335 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29336 = stablehlo.dynamic_slice %27307, %29335, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29337 = stablehlo.reshape %29336 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29338 = stablehlo.divide %29334, %29337 : tensor<f32> loc(#loc5253)
    %29339 = stablehlo.sqrt %29338 : tensor<f32> loc(#loc5249)
    %29340 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29341 = stablehlo.dynamic_slice %27323, %29340, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29342 = stablehlo.reshape %29341 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29343 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29344 = stablehlo.dynamic_slice %27323, %29343, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29345 = stablehlo.reshape %29344 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29346 = stablehlo.divide %29342, %29345 : tensor<f32> loc(#loc5253)
    %29347 = stablehlo.sqrt %29346 : tensor<f32> loc(#loc5249)
    %29348 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29349 = stablehlo.dynamic_slice %27339, %29348, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29350 = stablehlo.reshape %29349 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29351 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29352 = stablehlo.dynamic_slice %27339, %29351, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29353 = stablehlo.reshape %29352 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29354 = stablehlo.divide %29350, %29353 : tensor<f32> loc(#loc5253)
    %29355 = stablehlo.sqrt %29354 : tensor<f32> loc(#loc5249)
    %29356 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29357 = stablehlo.dynamic_slice %27355, %29356, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29358 = stablehlo.reshape %29357 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29359 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29360 = stablehlo.dynamic_slice %27355, %29359, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29361 = stablehlo.reshape %29360 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29362 = stablehlo.divide %29358, %29361 : tensor<f32> loc(#loc5253)
    %29363 = stablehlo.sqrt %29362 : tensor<f32> loc(#loc5249)
    %29364 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29365 = stablehlo.dynamic_slice %27371, %29364, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29366 = stablehlo.reshape %29365 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29367 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29368 = stablehlo.dynamic_slice %27371, %29367, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29369 = stablehlo.reshape %29368 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29370 = stablehlo.divide %29366, %29369 : tensor<f32> loc(#loc5253)
    %29371 = stablehlo.sqrt %29370 : tensor<f32> loc(#loc5249)
    %29372 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29373 = stablehlo.dynamic_slice %27387, %29372, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29374 = stablehlo.reshape %29373 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29375 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29376 = stablehlo.dynamic_slice %27387, %29375, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29377 = stablehlo.reshape %29376 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29378 = stablehlo.divide %29374, %29377 : tensor<f32> loc(#loc5253)
    %29379 = stablehlo.sqrt %29378 : tensor<f32> loc(#loc5249)
    %29380 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29381 = stablehlo.dynamic_slice %27403, %29380, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29382 = stablehlo.reshape %29381 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29383 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29384 = stablehlo.dynamic_slice %27403, %29383, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29385 = stablehlo.reshape %29384 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29386 = stablehlo.divide %29382, %29385 : tensor<f32> loc(#loc5253)
    %29387 = stablehlo.sqrt %29386 : tensor<f32> loc(#loc5249)
    %29388 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29389 = stablehlo.dynamic_slice %27419, %29388, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29390 = stablehlo.reshape %29389 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29391 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29392 = stablehlo.dynamic_slice %27419, %29391, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29393 = stablehlo.reshape %29392 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29394 = stablehlo.divide %29390, %29393 : tensor<f32> loc(#loc5253)
    %29395 = stablehlo.sqrt %29394 : tensor<f32> loc(#loc5249)
    %29396 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29397 = stablehlo.dynamic_slice %27435, %29396, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29398 = stablehlo.reshape %29397 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29399 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29400 = stablehlo.dynamic_slice %27435, %29399, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29401 = stablehlo.reshape %29400 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29402 = stablehlo.divide %29398, %29401 : tensor<f32> loc(#loc5253)
    %29403 = stablehlo.sqrt %29402 : tensor<f32> loc(#loc5249)
    %29404 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29405 = stablehlo.dynamic_slice %27451, %29404, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29406 = stablehlo.reshape %29405 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29407 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29408 = stablehlo.dynamic_slice %27451, %29407, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29409 = stablehlo.reshape %29408 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29410 = stablehlo.divide %29406, %29409 : tensor<f32> loc(#loc5253)
    %29411 = stablehlo.sqrt %29410 : tensor<f32> loc(#loc5249)
    %29412 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29413 = stablehlo.dynamic_slice %27467, %29412, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29414 = stablehlo.reshape %29413 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29415 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29416 = stablehlo.dynamic_slice %27467, %29415, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29417 = stablehlo.reshape %29416 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29418 = stablehlo.divide %29414, %29417 : tensor<f32> loc(#loc5253)
    %29419 = stablehlo.sqrt %29418 : tensor<f32> loc(#loc5249)
    %29420 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29421 = stablehlo.dynamic_slice %27483, %29420, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29422 = stablehlo.reshape %29421 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29423 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29424 = stablehlo.dynamic_slice %27483, %29423, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29425 = stablehlo.reshape %29424 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29426 = stablehlo.divide %29422, %29425 : tensor<f32> loc(#loc5253)
    %29427 = stablehlo.sqrt %29426 : tensor<f32> loc(#loc5249)
    %29428 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29429 = stablehlo.dynamic_slice %27499, %29428, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29430 = stablehlo.reshape %29429 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29431 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29432 = stablehlo.dynamic_slice %27499, %29431, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29433 = stablehlo.reshape %29432 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29434 = stablehlo.divide %29430, %29433 : tensor<f32> loc(#loc5253)
    %29435 = stablehlo.sqrt %29434 : tensor<f32> loc(#loc5249)
    %29436 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29437 = stablehlo.dynamic_slice %27515, %29436, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29438 = stablehlo.reshape %29437 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29439 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29440 = stablehlo.dynamic_slice %27515, %29439, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29441 = stablehlo.reshape %29440 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29442 = stablehlo.divide %29438, %29441 : tensor<f32> loc(#loc5253)
    %29443 = stablehlo.sqrt %29442 : tensor<f32> loc(#loc5249)
    %29444 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29445 = stablehlo.dynamic_slice %27531, %29444, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29446 = stablehlo.reshape %29445 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29447 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29448 = stablehlo.dynamic_slice %27531, %29447, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29449 = stablehlo.reshape %29448 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29450 = stablehlo.divide %29446, %29449 : tensor<f32> loc(#loc5253)
    %29451 = stablehlo.sqrt %29450 : tensor<f32> loc(#loc5249)
    %29452 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29453 = stablehlo.dynamic_slice %27547, %29452, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29454 = stablehlo.reshape %29453 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29455 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29456 = stablehlo.dynamic_slice %27547, %29455, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29457 = stablehlo.reshape %29456 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29458 = stablehlo.divide %29454, %29457 : tensor<f32> loc(#loc5253)
    %29459 = stablehlo.sqrt %29458 : tensor<f32> loc(#loc5249)
    %29460 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29461 = stablehlo.dynamic_slice %27563, %29460, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29462 = stablehlo.reshape %29461 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29463 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29464 = stablehlo.dynamic_slice %27563, %29463, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29465 = stablehlo.reshape %29464 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29466 = stablehlo.divide %29462, %29465 : tensor<f32> loc(#loc5253)
    %29467 = stablehlo.sqrt %29466 : tensor<f32> loc(#loc5249)
    %29468 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29469 = stablehlo.dynamic_slice %27579, %29468, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29470 = stablehlo.reshape %29469 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29471 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29472 = stablehlo.dynamic_slice %27579, %29471, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29473 = stablehlo.reshape %29472 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29474 = stablehlo.divide %29470, %29473 : tensor<f32> loc(#loc5253)
    %29475 = stablehlo.sqrt %29474 : tensor<f32> loc(#loc5249)
    %29476 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29477 = stablehlo.dynamic_slice %27595, %29476, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29478 = stablehlo.reshape %29477 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29479 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29480 = stablehlo.dynamic_slice %27595, %29479, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29481 = stablehlo.reshape %29480 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29482 = stablehlo.divide %29478, %29481 : tensor<f32> loc(#loc5253)
    %29483 = stablehlo.sqrt %29482 : tensor<f32> loc(#loc5249)
    %29484 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29485 = stablehlo.dynamic_slice %27611, %29484, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29486 = stablehlo.reshape %29485 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29487 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29488 = stablehlo.dynamic_slice %27611, %29487, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29489 = stablehlo.reshape %29488 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29490 = stablehlo.divide %29486, %29489 : tensor<f32> loc(#loc5253)
    %29491 = stablehlo.sqrt %29490 : tensor<f32> loc(#loc5249)
    %29492 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29493 = stablehlo.dynamic_slice %27627, %29492, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29494 = stablehlo.reshape %29493 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29495 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29496 = stablehlo.dynamic_slice %27627, %29495, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29497 = stablehlo.reshape %29496 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29498 = stablehlo.divide %29494, %29497 : tensor<f32> loc(#loc5253)
    %29499 = stablehlo.sqrt %29498 : tensor<f32> loc(#loc5249)
    %29500 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29501 = stablehlo.dynamic_slice %27643, %29500, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29502 = stablehlo.reshape %29501 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29503 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29504 = stablehlo.dynamic_slice %27643, %29503, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29505 = stablehlo.reshape %29504 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29506 = stablehlo.divide %29502, %29505 : tensor<f32> loc(#loc5253)
    %29507 = stablehlo.sqrt %29506 : tensor<f32> loc(#loc5249)
    %29508 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29509 = stablehlo.dynamic_slice %27659, %29508, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29510 = stablehlo.reshape %29509 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29511 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29512 = stablehlo.dynamic_slice %27659, %29511, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29513 = stablehlo.reshape %29512 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29514 = stablehlo.divide %29510, %29513 : tensor<f32> loc(#loc5253)
    %29515 = stablehlo.sqrt %29514 : tensor<f32> loc(#loc5249)
    %29516 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29517 = stablehlo.dynamic_slice %27675, %29516, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29518 = stablehlo.reshape %29517 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29519 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29520 = stablehlo.dynamic_slice %27675, %29519, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29521 = stablehlo.reshape %29520 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29522 = stablehlo.divide %29518, %29521 : tensor<f32> loc(#loc5253)
    %29523 = stablehlo.sqrt %29522 : tensor<f32> loc(#loc5249)
    %29524 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29525 = stablehlo.dynamic_slice %27691, %29524, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29526 = stablehlo.reshape %29525 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29527 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29528 = stablehlo.dynamic_slice %27691, %29527, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29529 = stablehlo.reshape %29528 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29530 = stablehlo.divide %29526, %29529 : tensor<f32> loc(#loc5253)
    %29531 = stablehlo.sqrt %29530 : tensor<f32> loc(#loc5249)
    %29532 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29533 = stablehlo.dynamic_slice %27707, %29532, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29534 = stablehlo.reshape %29533 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29535 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29536 = stablehlo.dynamic_slice %27707, %29535, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29537 = stablehlo.reshape %29536 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29538 = stablehlo.divide %29534, %29537 : tensor<f32> loc(#loc5253)
    %29539 = stablehlo.sqrt %29538 : tensor<f32> loc(#loc5249)
    %29540 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29541 = stablehlo.dynamic_slice %27723, %29540, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29542 = stablehlo.reshape %29541 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29543 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29544 = stablehlo.dynamic_slice %27723, %29543, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29545 = stablehlo.reshape %29544 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29546 = stablehlo.divide %29542, %29545 : tensor<f32> loc(#loc5253)
    %29547 = stablehlo.sqrt %29546 : tensor<f32> loc(#loc5249)
    %29548 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29549 = stablehlo.dynamic_slice %27739, %29548, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29550 = stablehlo.reshape %29549 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29551 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29552 = stablehlo.dynamic_slice %27739, %29551, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29553 = stablehlo.reshape %29552 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29554 = stablehlo.divide %29550, %29553 : tensor<f32> loc(#loc5253)
    %29555 = stablehlo.sqrt %29554 : tensor<f32> loc(#loc5249)
    %29556 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29557 = stablehlo.dynamic_slice %27755, %29556, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29558 = stablehlo.reshape %29557 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29559 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29560 = stablehlo.dynamic_slice %27755, %29559, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29561 = stablehlo.reshape %29560 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29562 = stablehlo.divide %29558, %29561 : tensor<f32> loc(#loc5253)
    %29563 = stablehlo.sqrt %29562 : tensor<f32> loc(#loc5249)
    %29564 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29565 = stablehlo.dynamic_slice %27771, %29564, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29566 = stablehlo.reshape %29565 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29567 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29568 = stablehlo.dynamic_slice %27771, %29567, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29569 = stablehlo.reshape %29568 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29570 = stablehlo.divide %29566, %29569 : tensor<f32> loc(#loc5253)
    %29571 = stablehlo.sqrt %29570 : tensor<f32> loc(#loc5249)
    %29572 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29573 = stablehlo.dynamic_slice %27787, %29572, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29574 = stablehlo.reshape %29573 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29575 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29576 = stablehlo.dynamic_slice %27787, %29575, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29577 = stablehlo.reshape %29576 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29578 = stablehlo.divide %29574, %29577 : tensor<f32> loc(#loc5253)
    %29579 = stablehlo.sqrt %29578 : tensor<f32> loc(#loc5249)
    %29580 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29581 = stablehlo.dynamic_slice %27803, %29580, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29582 = stablehlo.reshape %29581 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29583 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29584 = stablehlo.dynamic_slice %27803, %29583, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29585 = stablehlo.reshape %29584 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29586 = stablehlo.divide %29582, %29585 : tensor<f32> loc(#loc5253)
    %29587 = stablehlo.sqrt %29586 : tensor<f32> loc(#loc5249)
    %29588 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29589 = stablehlo.dynamic_slice %27819, %29588, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29590 = stablehlo.reshape %29589 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29591 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29592 = stablehlo.dynamic_slice %27819, %29591, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29593 = stablehlo.reshape %29592 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29594 = stablehlo.divide %29590, %29593 : tensor<f32> loc(#loc5253)
    %29595 = stablehlo.sqrt %29594 : tensor<f32> loc(#loc5249)
    %29596 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29597 = stablehlo.dynamic_slice %27835, %29596, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29598 = stablehlo.reshape %29597 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29599 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29600 = stablehlo.dynamic_slice %27835, %29599, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29601 = stablehlo.reshape %29600 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29602 = stablehlo.divide %29598, %29601 : tensor<f32> loc(#loc5253)
    %29603 = stablehlo.sqrt %29602 : tensor<f32> loc(#loc5249)
    %29604 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29605 = stablehlo.dynamic_slice %27851, %29604, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29606 = stablehlo.reshape %29605 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29607 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29608 = stablehlo.dynamic_slice %27851, %29607, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29609 = stablehlo.reshape %29608 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29610 = stablehlo.divide %29606, %29609 : tensor<f32> loc(#loc5253)
    %29611 = stablehlo.sqrt %29610 : tensor<f32> loc(#loc5249)
    %29612 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29613 = stablehlo.dynamic_slice %27867, %29612, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29614 = stablehlo.reshape %29613 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29615 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29616 = stablehlo.dynamic_slice %27867, %29615, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29617 = stablehlo.reshape %29616 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29618 = stablehlo.divide %29614, %29617 : tensor<f32> loc(#loc5253)
    %29619 = stablehlo.sqrt %29618 : tensor<f32> loc(#loc5249)
    %29620 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29621 = stablehlo.dynamic_slice %27883, %29620, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29622 = stablehlo.reshape %29621 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29623 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29624 = stablehlo.dynamic_slice %27883, %29623, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29625 = stablehlo.reshape %29624 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29626 = stablehlo.divide %29622, %29625 : tensor<f32> loc(#loc5253)
    %29627 = stablehlo.sqrt %29626 : tensor<f32> loc(#loc5249)
    %29628 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29629 = stablehlo.dynamic_slice %27899, %29628, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29630 = stablehlo.reshape %29629 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29631 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29632 = stablehlo.dynamic_slice %27899, %29631, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29633 = stablehlo.reshape %29632 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29634 = stablehlo.divide %29630, %29633 : tensor<f32> loc(#loc5253)
    %29635 = stablehlo.sqrt %29634 : tensor<f32> loc(#loc5249)
    %29636 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29637 = stablehlo.dynamic_slice %27915, %29636, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29638 = stablehlo.reshape %29637 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29639 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29640 = stablehlo.dynamic_slice %27915, %29639, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29641 = stablehlo.reshape %29640 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29642 = stablehlo.divide %29638, %29641 : tensor<f32> loc(#loc5253)
    %29643 = stablehlo.sqrt %29642 : tensor<f32> loc(#loc5249)
    %29644 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29645 = stablehlo.dynamic_slice %27931, %29644, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29646 = stablehlo.reshape %29645 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29647 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29648 = stablehlo.dynamic_slice %27931, %29647, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29649 = stablehlo.reshape %29648 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29650 = stablehlo.divide %29646, %29649 : tensor<f32> loc(#loc5253)
    %29651 = stablehlo.sqrt %29650 : tensor<f32> loc(#loc5249)
    %29652 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29653 = stablehlo.dynamic_slice %27947, %29652, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29654 = stablehlo.reshape %29653 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29655 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29656 = stablehlo.dynamic_slice %27947, %29655, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29657 = stablehlo.reshape %29656 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29658 = stablehlo.divide %29654, %29657 : tensor<f32> loc(#loc5253)
    %29659 = stablehlo.sqrt %29658 : tensor<f32> loc(#loc5249)
    %29660 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29661 = stablehlo.dynamic_slice %27963, %29660, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29662 = stablehlo.reshape %29661 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29663 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29664 = stablehlo.dynamic_slice %27963, %29663, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29665 = stablehlo.reshape %29664 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29666 = stablehlo.divide %29662, %29665 : tensor<f32> loc(#loc5253)
    %29667 = stablehlo.sqrt %29666 : tensor<f32> loc(#loc5249)
    %29668 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29669 = stablehlo.dynamic_slice %27979, %29668, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29670 = stablehlo.reshape %29669 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29671 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29672 = stablehlo.dynamic_slice %27979, %29671, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29673 = stablehlo.reshape %29672 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29674 = stablehlo.divide %29670, %29673 : tensor<f32> loc(#loc5253)
    %29675 = stablehlo.sqrt %29674 : tensor<f32> loc(#loc5249)
    %29676 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29677 = stablehlo.dynamic_slice %27995, %29676, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29678 = stablehlo.reshape %29677 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29679 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29680 = stablehlo.dynamic_slice %27995, %29679, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29681 = stablehlo.reshape %29680 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29682 = stablehlo.divide %29678, %29681 : tensor<f32> loc(#loc5253)
    %29683 = stablehlo.sqrt %29682 : tensor<f32> loc(#loc5249)
    %29684 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29685 = stablehlo.dynamic_slice %28011, %29684, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29686 = stablehlo.reshape %29685 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29687 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29688 = stablehlo.dynamic_slice %28011, %29687, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29689 = stablehlo.reshape %29688 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29690 = stablehlo.divide %29686, %29689 : tensor<f32> loc(#loc5253)
    %29691 = stablehlo.sqrt %29690 : tensor<f32> loc(#loc5249)
    %29692 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29693 = stablehlo.dynamic_slice %28027, %29692, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29694 = stablehlo.reshape %29693 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29695 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29696 = stablehlo.dynamic_slice %28027, %29695, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29697 = stablehlo.reshape %29696 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29698 = stablehlo.divide %29694, %29697 : tensor<f32> loc(#loc5253)
    %29699 = stablehlo.sqrt %29698 : tensor<f32> loc(#loc5249)
    %29700 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29701 = stablehlo.dynamic_slice %28043, %29700, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29702 = stablehlo.reshape %29701 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29703 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29704 = stablehlo.dynamic_slice %28043, %29703, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29705 = stablehlo.reshape %29704 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29706 = stablehlo.divide %29702, %29705 : tensor<f32> loc(#loc5253)
    %29707 = stablehlo.sqrt %29706 : tensor<f32> loc(#loc5249)
    %29708 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29709 = stablehlo.dynamic_slice %28059, %29708, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29710 = stablehlo.reshape %29709 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29711 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29712 = stablehlo.dynamic_slice %28059, %29711, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29713 = stablehlo.reshape %29712 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29714 = stablehlo.divide %29710, %29713 : tensor<f32> loc(#loc5253)
    %29715 = stablehlo.sqrt %29714 : tensor<f32> loc(#loc5249)
    %29716 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29717 = stablehlo.dynamic_slice %28075, %29716, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29718 = stablehlo.reshape %29717 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29719 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29720 = stablehlo.dynamic_slice %28075, %29719, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29721 = stablehlo.reshape %29720 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29722 = stablehlo.divide %29718, %29721 : tensor<f32> loc(#loc5253)
    %29723 = stablehlo.sqrt %29722 : tensor<f32> loc(#loc5249)
    %29724 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29725 = stablehlo.dynamic_slice %28091, %29724, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29726 = stablehlo.reshape %29725 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29727 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29728 = stablehlo.dynamic_slice %28091, %29727, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29729 = stablehlo.reshape %29728 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29730 = stablehlo.divide %29726, %29729 : tensor<f32> loc(#loc5253)
    %29731 = stablehlo.sqrt %29730 : tensor<f32> loc(#loc5249)
    %29732 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29733 = stablehlo.dynamic_slice %28107, %29732, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29734 = stablehlo.reshape %29733 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29735 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29736 = stablehlo.dynamic_slice %28107, %29735, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29737 = stablehlo.reshape %29736 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29738 = stablehlo.divide %29734, %29737 : tensor<f32> loc(#loc5253)
    %29739 = stablehlo.sqrt %29738 : tensor<f32> loc(#loc5249)
    %29740 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29741 = stablehlo.dynamic_slice %28123, %29740, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29742 = stablehlo.reshape %29741 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29743 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29744 = stablehlo.dynamic_slice %28123, %29743, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29745 = stablehlo.reshape %29744 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29746 = stablehlo.divide %29742, %29745 : tensor<f32> loc(#loc5253)
    %29747 = stablehlo.sqrt %29746 : tensor<f32> loc(#loc5249)
    %29748 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29749 = stablehlo.dynamic_slice %28139, %29748, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29750 = stablehlo.reshape %29749 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29751 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29752 = stablehlo.dynamic_slice %28139, %29751, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29753 = stablehlo.reshape %29752 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29754 = stablehlo.divide %29750, %29753 : tensor<f32> loc(#loc5253)
    %29755 = stablehlo.sqrt %29754 : tensor<f32> loc(#loc5249)
    %29756 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29757 = stablehlo.dynamic_slice %28155, %29756, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29758 = stablehlo.reshape %29757 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29759 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29760 = stablehlo.dynamic_slice %28155, %29759, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29761 = stablehlo.reshape %29760 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29762 = stablehlo.divide %29758, %29761 : tensor<f32> loc(#loc5253)
    %29763 = stablehlo.sqrt %29762 : tensor<f32> loc(#loc5249)
    %29764 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29765 = stablehlo.dynamic_slice %28171, %29764, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29766 = stablehlo.reshape %29765 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29767 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29768 = stablehlo.dynamic_slice %28171, %29767, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29769 = stablehlo.reshape %29768 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29770 = stablehlo.divide %29766, %29769 : tensor<f32> loc(#loc5253)
    %29771 = stablehlo.sqrt %29770 : tensor<f32> loc(#loc5249)
    %29772 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29773 = stablehlo.dynamic_slice %28187, %29772, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29774 = stablehlo.reshape %29773 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29775 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29776 = stablehlo.dynamic_slice %28187, %29775, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29777 = stablehlo.reshape %29776 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29778 = stablehlo.divide %29774, %29777 : tensor<f32> loc(#loc5253)
    %29779 = stablehlo.sqrt %29778 : tensor<f32> loc(#loc5249)
    %29780 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29781 = stablehlo.dynamic_slice %28203, %29780, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29782 = stablehlo.reshape %29781 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29783 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29784 = stablehlo.dynamic_slice %28203, %29783, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29785 = stablehlo.reshape %29784 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29786 = stablehlo.divide %29782, %29785 : tensor<f32> loc(#loc5253)
    %29787 = stablehlo.sqrt %29786 : tensor<f32> loc(#loc5249)
    %29788 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29789 = stablehlo.dynamic_slice %28219, %29788, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29790 = stablehlo.reshape %29789 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29791 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29792 = stablehlo.dynamic_slice %28219, %29791, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29793 = stablehlo.reshape %29792 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29794 = stablehlo.divide %29790, %29793 : tensor<f32> loc(#loc5253)
    %29795 = stablehlo.sqrt %29794 : tensor<f32> loc(#loc5249)
    %29796 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29797 = stablehlo.dynamic_slice %28235, %29796, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29798 = stablehlo.reshape %29797 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29799 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29800 = stablehlo.dynamic_slice %28235, %29799, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29801 = stablehlo.reshape %29800 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29802 = stablehlo.divide %29798, %29801 : tensor<f32> loc(#loc5253)
    %29803 = stablehlo.sqrt %29802 : tensor<f32> loc(#loc5249)
    %29804 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29805 = stablehlo.dynamic_slice %28251, %29804, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29806 = stablehlo.reshape %29805 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29807 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29808 = stablehlo.dynamic_slice %28251, %29807, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29809 = stablehlo.reshape %29808 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29810 = stablehlo.divide %29806, %29809 : tensor<f32> loc(#loc5253)
    %29811 = stablehlo.sqrt %29810 : tensor<f32> loc(#loc5249)
    %29812 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29813 = stablehlo.dynamic_slice %28267, %29812, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29814 = stablehlo.reshape %29813 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29815 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29816 = stablehlo.dynamic_slice %28267, %29815, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29817 = stablehlo.reshape %29816 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29818 = stablehlo.divide %29814, %29817 : tensor<f32> loc(#loc5253)
    %29819 = stablehlo.sqrt %29818 : tensor<f32> loc(#loc5249)
    %29820 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29821 = stablehlo.dynamic_slice %28283, %29820, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29822 = stablehlo.reshape %29821 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29823 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29824 = stablehlo.dynamic_slice %28283, %29823, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29825 = stablehlo.reshape %29824 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29826 = stablehlo.divide %29822, %29825 : tensor<f32> loc(#loc5253)
    %29827 = stablehlo.sqrt %29826 : tensor<f32> loc(#loc5249)
    %29828 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29829 = stablehlo.dynamic_slice %28299, %29828, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29830 = stablehlo.reshape %29829 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29831 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29832 = stablehlo.dynamic_slice %28299, %29831, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29833 = stablehlo.reshape %29832 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29834 = stablehlo.divide %29830, %29833 : tensor<f32> loc(#loc5253)
    %29835 = stablehlo.sqrt %29834 : tensor<f32> loc(#loc5249)
    %29836 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29837 = stablehlo.dynamic_slice %28315, %29836, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29838 = stablehlo.reshape %29837 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29839 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29840 = stablehlo.dynamic_slice %28315, %29839, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29841 = stablehlo.reshape %29840 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29842 = stablehlo.divide %29838, %29841 : tensor<f32> loc(#loc5253)
    %29843 = stablehlo.sqrt %29842 : tensor<f32> loc(#loc5249)
    %29844 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29845 = stablehlo.dynamic_slice %28331, %29844, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29846 = stablehlo.reshape %29845 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29847 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29848 = stablehlo.dynamic_slice %28331, %29847, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29849 = stablehlo.reshape %29848 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29850 = stablehlo.divide %29846, %29849 : tensor<f32> loc(#loc5253)
    %29851 = stablehlo.sqrt %29850 : tensor<f32> loc(#loc5249)
    %29852 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29853 = stablehlo.dynamic_slice %28347, %29852, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29854 = stablehlo.reshape %29853 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29855 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29856 = stablehlo.dynamic_slice %28347, %29855, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29857 = stablehlo.reshape %29856 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29858 = stablehlo.divide %29854, %29857 : tensor<f32> loc(#loc5253)
    %29859 = stablehlo.sqrt %29858 : tensor<f32> loc(#loc5249)
    %29860 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29861 = stablehlo.dynamic_slice %28363, %29860, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29862 = stablehlo.reshape %29861 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29863 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29864 = stablehlo.dynamic_slice %28363, %29863, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29865 = stablehlo.reshape %29864 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29866 = stablehlo.divide %29862, %29865 : tensor<f32> loc(#loc5253)
    %29867 = stablehlo.sqrt %29866 : tensor<f32> loc(#loc5249)
    %29868 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29869 = stablehlo.dynamic_slice %28379, %29868, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29870 = stablehlo.reshape %29869 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29871 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29872 = stablehlo.dynamic_slice %28379, %29871, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29873 = stablehlo.reshape %29872 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29874 = stablehlo.divide %29870, %29873 : tensor<f32> loc(#loc5253)
    %29875 = stablehlo.sqrt %29874 : tensor<f32> loc(#loc5249)
    %29876 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29877 = stablehlo.dynamic_slice %28395, %29876, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29878 = stablehlo.reshape %29877 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29879 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29880 = stablehlo.dynamic_slice %28395, %29879, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29881 = stablehlo.reshape %29880 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29882 = stablehlo.divide %29878, %29881 : tensor<f32> loc(#loc5253)
    %29883 = stablehlo.sqrt %29882 : tensor<f32> loc(#loc5249)
    %29884 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29885 = stablehlo.dynamic_slice %28411, %29884, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29886 = stablehlo.reshape %29885 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29887 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29888 = stablehlo.dynamic_slice %28411, %29887, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29889 = stablehlo.reshape %29888 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29890 = stablehlo.divide %29886, %29889 : tensor<f32> loc(#loc5253)
    %29891 = stablehlo.sqrt %29890 : tensor<f32> loc(#loc5249)
    %29892 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29893 = stablehlo.dynamic_slice %28427, %29892, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29894 = stablehlo.reshape %29893 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29895 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29896 = stablehlo.dynamic_slice %28427, %29895, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29897 = stablehlo.reshape %29896 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29898 = stablehlo.divide %29894, %29897 : tensor<f32> loc(#loc5253)
    %29899 = stablehlo.sqrt %29898 : tensor<f32> loc(#loc5249)
    %29900 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29901 = stablehlo.dynamic_slice %28443, %29900, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29902 = stablehlo.reshape %29901 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29903 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29904 = stablehlo.dynamic_slice %28443, %29903, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29905 = stablehlo.reshape %29904 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29906 = stablehlo.divide %29902, %29905 : tensor<f32> loc(#loc5253)
    %29907 = stablehlo.sqrt %29906 : tensor<f32> loc(#loc5249)
    %29908 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29909 = stablehlo.dynamic_slice %28459, %29908, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29910 = stablehlo.reshape %29909 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29911 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29912 = stablehlo.dynamic_slice %28459, %29911, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29913 = stablehlo.reshape %29912 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29914 = stablehlo.divide %29910, %29913 : tensor<f32> loc(#loc5253)
    %29915 = stablehlo.sqrt %29914 : tensor<f32> loc(#loc5249)
    %29916 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %29917 = stablehlo.dynamic_slice %28475, %29916, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29918 = stablehlo.reshape %29917 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29919 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %29920 = stablehlo.dynamic_slice %28475, %29919, sizes = [1] : (tensor<2xf32>, tensor<i32>) -> tensor<1xf32> loc(#loc5318)
    %29921 = stablehlo.reshape %29920 : (tensor<1xf32>) -> tensor<f32> loc(#loc5319)
    %29922 = stablehlo.divide %29918, %29921 : tensor<f32> loc(#loc5253)
    %29923 = stablehlo.sqrt %29922 : tensor<f32> loc(#loc5249)
    %29924 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %29925 = stablehlo.constant dense<1.000000e+00> : tensor<bf16> loc(#loc)
    %29926 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %29927 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    %29928 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc)
    return %25579, %25397, %25398, %25399, %25400, %25401, %25402, %25403, %25404, %25405, %25406, %25407, %25408, %25409, %25410, %25411, %25412, %25413, %25414, %25415, %25416, %25417, %25418, %25419, %25420, %25421, %25422, %25423, %25424, %25425, %25426, %25427, %25428, %25429, %25430, %25431, %25432, %25433, %25434, %25435, %25436, %25437, %25438, %25439, %25440, %25441, %25442, %25443, %25444, %25445, %25446, %25447, %25448, %25449, %25450, %25451, %25452, %25453, %25454, %25455, %25456, %25457, %25458, %25459, %25460, %25461, %25462, %25463, %25464, %25465, %25466, %25467, %25468, %25469, %25470, %25471, %25472, %25473, %25474, %25475, %25476, %25477, %25478, %25479, %25480, %25481, %25482, %25483, %25484, %25485, %25486, %25487, %25488, %25489, %25490, %25491, %25492, %25493, %25494, %25495, %25496, %25497, %25498, %25499, %25500, %25501, %25502, %25503, %25504, %25505, %25506, %25507, %25508, %25509, %25510, %25511, %25512, %25513, %25514, %25515, %25516, %25517, %25518, %25519, %25520, %25521, %25522, %25523, %25524, %25525, %25526, %25527, %25528, %25529, %25530, %25531, %25532, %25533, %25534, %25535, %25536, %25537, %25538, %25539, %25540, %25541, %25542, %25543, %25544, %25545, %25546, %25547, %25548, %25549, %25550, %25551, %25552, %25553, %25554, %25555, %25556, %25557, %25558, %25559, %25560, %25561, %25562, %25563, %25564, %25565, %25566, %25567, %25568, %25569, %25570, %25571, %25572, %25573, %25574, %25575, %25576, %25577, %23547, %23548, %23549, %23550, %23551, %23552, %23553, %23554, %23555, %23556, %23557, %23558, %23559, %23560, %23561, %23562, %23563, %23564, %23565, %23566, %23567, %23568, %23569, %23570, %23571, %23572, %23573, %23574, %23575, %23576, %23577, %23578, %23579, %23580, %23581, %23582, %23583, %23584, %23585, %23586, %23587, %23588, %23589, %23590, %23591, %23592, %23593, %23594, %23595, %23596, %23597, %23598, %23599, %23600, %23601, %23602, %23603, %23604, %23605, %23606, %23607, %23608, %23609, %23610, %23611, %23612, %23613, %23614, %23615, %23616, %23617, %23618, %23619, %23620, %23621, %23622, %23623, %23624, %23625, %23626, %23627, %23628, %23629, %23630, %23631, %23632, %23633, %23634, %23635, %23636, %23637, %23638, %23639, %23640, %23641, %23642, %23643, %23644, %23645, %23646, %23647, %23648, %23649, %23650, %23651, %23652, %23653, %23654, %23655, %23656, %23657, %23658, %23659, %23660, %23661, %23662, %23663, %23664, %23665, %23666, %23667, %23668, %23669, %23670, %23671, %23672, %23673, %23674, %23675, %23676, %23677, %23678, %23679, %23680, %23681, %23682, %23683, %23684, %23685, %23686, %23687, %23688, %23689, %23690, %23691, %23692, %23693, %23694, %23695, %23696, %23697, %23698, %23699, %23700, %23701, %23702, %23703, %23704, %23705, %23706, %23707, %23708, %23709, %23710, %23711, %23712, %23713, %23714, %23715, %23716, %23717, %23718, %23719, %23720, %23721, %23722, %23723, %23724, %23725, %23726, %23727, %23728, %23729, %23730, %23731, %23732, %23733, %23734, %23735, %23736, %23737, %23738, %23739, %23740, %23741, %23742, %23743, %23744, %23745, %23746, %23747, %23748, %23749, %23750, %23751, %23752, %23753, %23754, %23755, %23756, %23757, %23758, %23759, %23760, %23761, %23762, %23763, %23764, %23765, %23766, %23767, %23768, %23769, %23770, %23771, %23772, %23773, %23774, %23775, %23776, %23777, %23778, %23779, %23780, %23781, %23782, %23783, %23784, %23785, %23786, %23787, %23788, %23789, %23790, %23791, %23792, %23793, %23794, %23795, %23796, %23797, %23798, %23799, %23800, %23801, %23802, %23803, %23804, %23805, %23806, %23807, %23808, %23809, %23810, %23811, %23812, %23813, %23814, %23815, %23816, %23817, %23818, %23819, %23820, %23821, %23822, %23823, %23824, %23825, %23826, %23827, %23828, %23829, %23830, %23831, %23832, %23833, %23834, %23835, %23836, %23837, %23838, %23839, %23840, %23841, %23842, %23843, %23844, %23845, %23846, %23847, %23848, %23849, %23850, %23851, %23852, %23853, %23854, %23855, %23856, %23857, %23858, %23859, %23860, %23861, %23862, %23863, %23864, %23865, %23866, %23867, %23868, %23869, %23870, %23871, %23872, %23873, %23874, %23875, %23876, %23877, %23878, %23879, %23880, %23881, %23882, %23883, %23884, %23885, %23886, %23887, %23888, %23889, %23890, %23891, %23892, %23893, %23894, %23895, %23896, %23897, %23898, %23899, %23900, %23901, %23902, %23903, %23904, %23905, %23906, %23907, %23908, %23909, %23910, %23911, %23912, %23913, %23914, %23915, %23916, %1153, %29924, %29925, %1147, %1144, %1162, %1144, %1147, %1144, %1144, %29926, %1153, %1144, %7, %1163, %24656, %3801, %2699, %2695, %1951, %2691, %25396, %198, %200, %29927, %29928, %919, %922, %925, %928, %931, %934, %937, %940, %943, %946, %949, %952, %955, %958, %961, %964, %967, %970, %973, %976, %979, %982, %985, %988, %991, %994, %997, %1000, %1003, %1006, %1009, %1012, %1015, %1018, %1021, %1024, %1027, %1030, %1033, %1036, %1039, %1042, %1045, %1048, %1051, %1054, %1057, %1060, %1951, %28483, %28491, %28499, %28507, %28515, %28523, %28531, %28539, %28547, %28555, %28563, %28571, %28579, %28587, %28595, %28603, %28611, %28619, %28627, %28635, %28643, %28651, %28659, %28667, %28675, %28683, %28691, %28699, %28707, %28715, %28723, %28731, %28739, %28747, %28755, %28763, %28771, %28779, %28787, %28795, %28803, %28811, %28819, %28827, %28835, %28843, %28851, %28859, %28867, %28875, %28883, %28891, %28899, %28907, %28915, %28923, %28931, %28939, %28947, %28955, %28963, %28971, %28979, %28987, %28995, %29003, %29011, %29019, %29027, %29035, %29043, %29051, %29059, %29067, %29075, %29083, %29091, %29099, %29107, %29115, %29123, %29131, %29139, %29147, %29155, %29163, %29171, %29179, %29187, %29195, %29203, %29211, %29219, %29227, %29235, %29243, %29251, %29259, %29267, %29275, %29283, %29291, %29299, %29307, %29315, %29323, %29331, %29339, %29347, %29355, %29363, %29371, %29379, %29387, %29395, %29403, %29411, %29419, %29427, %29435, %29443, %29451, %29459, %29467, %29475, %29483, %29491, %29499, %29507, %29515, %29523, %29531, %29539, %29547, %29555, %29563, %29571, %29579, %29587, %29595, %29603, %29611, %29619, %29627, %29635, %29643, %29651, %29659, %29667, %29675, %29683, %29691, %29699, %29707, %29715, %29723, %29731, %29739, %29747, %29755, %29763, %29771, %29779, %29787, %29795, %29803, %29811, %29819, %29827, %29835, %29843, %29851, %29859, %29867, %29875, %29883, %29891, %29899, %29907, %29915, %29923 : tensor<ui32>, tensor<2048xf32>, tensor<2048xf32>, tensor<2048x2048xf32>, tensor<51200xf32>, tensor<2048x51200xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<2048xf32>, tensor<2048xf32>, tensor<2048x2048xf32>, tensor<51200xf32>, tensor<2048x51200xf32>, tensor<2048xf32>, tensor<2048xf32>, tensor<2048x2048xf32>, tensor<51200xf32>, tensor<2048x51200xf32>, tensor<i32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1x8192xf32>, tensor<1x2048x8192xf32>, tensor<1x2048xf32>, tensor<1x8192x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x2048xf32>, tensor<1x3x2048x32x64xf32>, tensor<1x64xf32>, tensor<1x2048x32x64xf32>, tensor<1xi32>, tensor<f32>, tensor<bf16>, tensor<bf16>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<1x2048xi32>, tensor<1xf32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<bf16>, tensor<bf16>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32> loc(#loc)
  } loc(#loc)
  func.func private @_one_hot(%arg0: tensor<1x2048xi32> loc(unknown)) -> tensor<1x2048x51200xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<1x2048xi32>) -> tensor<1x2048x1xi32> loc(#loc5320)
    %1 = stablehlo.iota dim = 2 : tensor<1x1x51200xi32> loc(#loc5321)
    %2 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 2] : (tensor<1x2048x1xi32>) -> tensor<1x2048x51200xi32> loc(#loc5322)
    %3 = stablehlo.broadcast_in_dim %1, dims = [0, 1, 2] : (tensor<1x1x51200xi32>) -> tensor<1x2048x51200xi32> loc(#loc5322)
    %4 = stablehlo.compare  EQ, %2, %3,  SIGNED : (tensor<1x2048x51200xi32>, tensor<1x2048x51200xi32>) -> tensor<1x2048x51200xi1> loc(#loc5322)
    %5 = stablehlo.convert %4 : (tensor<1x2048x51200xi1>) -> tensor<1x2048x51200xbf16> loc(#loc5323)
    return %5 : tensor<1x2048x51200xbf16> loc(#loc141)
  } loc(#loc141)
  func.func private @_einsum(%arg0: tensor<1x2048x51200xbf16> loc(unknown), %arg1: tensor<51200x2048xbf16> loc(unknown)) -> tensor<1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x51200xbf16>, tensor<51200x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc5324)
    return %0 : tensor<1x2048x2048xbf16> loc(#loc142)
  } loc(#loc142)
  func.func private @_one_hot_0(%arg0: tensor<1x2048xi32> loc(unknown)) -> tensor<1x2048x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<1x2048xi32>) -> tensor<1x2048x1xi32> loc(#loc5325)
    %1 = stablehlo.iota dim = 2 : tensor<1x1x2048xi32> loc(#loc5326)
    %2 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 2] : (tensor<1x2048x1xi32>) -> tensor<1x2048x2048xi32> loc(#loc5327)
    %3 = stablehlo.broadcast_in_dim %1, dims = [0, 1, 2] : (tensor<1x1x2048xi32>) -> tensor<1x2048x2048xi32> loc(#loc5327)
    %4 = stablehlo.compare  EQ, %2, %3,  SIGNED : (tensor<1x2048x2048xi32>, tensor<1x2048x2048xi32>) -> tensor<1x2048x2048xi1> loc(#loc5327)
    %5 = stablehlo.convert %4 : (tensor<1x2048x2048xi1>) -> tensor<1x2048x2048xbf16> loc(#loc5328)
    return %5 : tensor<1x2048x2048xbf16> loc(#loc146)
  } loc(#loc146)
  func.func private @_einsum_1(%arg0: tensor<1x2048x2048xbf16> loc(unknown), %arg1: tensor<2048x2048xbf16> loc(unknown)) -> tensor<1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x2048xbf16>, tensor<2048x2048xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc5329)
    return %0 : tensor<1x2048x2048xbf16> loc(#loc147)
  } loc(#loc147)
  func.func private @_pad(%arg0: tensor<1x1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<i32> loc(unknown)) -> tensor<1x1x1x2048x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<i32>) -> tensor<5x2xi32> loc(#loc5330)
    %1 = stablehlo.convert %0 : (tensor<5x2xi32>) -> tensor<5x2xbf16> loc(#loc5331)
    %2 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %4 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %5 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %6 = stablehlo.concatenate %3, %5, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %7 = "stablehlo.gather"(%1, %6) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %8 = stablehlo.pad %arg0, %7, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %9 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %11 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %12 = stablehlo.broadcast_in_dim %11, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %13 = stablehlo.concatenate %10, %12, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %14 = "stablehlo.gather"(%1, %13) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %15 = stablehlo.pad %8, %14, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %16 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %17 = stablehlo.broadcast_in_dim %16, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %18 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %19 = stablehlo.broadcast_in_dim %18, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %20 = stablehlo.concatenate %17, %19, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %21 = "stablehlo.gather"(%1, %20) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %22 = stablehlo.pad %15, %21, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %23 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %24 = stablehlo.broadcast_in_dim %23, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %25 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %26 = stablehlo.broadcast_in_dim %25, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %27 = stablehlo.concatenate %24, %26, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %28 = "stablehlo.gather"(%1, %27) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %29 = stablehlo.pad %22, %28, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %30 = stablehlo.constant dense<2> : tensor<i32> loc(#loc185)
    %31 = stablehlo.broadcast_in_dim %30, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %32 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %33 = stablehlo.broadcast_in_dim %32, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %34 = stablehlo.concatenate %31, %33, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %35 = "stablehlo.gather"(%1, %34) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %36 = stablehlo.pad %29, %35, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %37 = stablehlo.constant dense<2> : tensor<i32> loc(#loc185)
    %38 = stablehlo.broadcast_in_dim %37, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %39 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %40 = stablehlo.broadcast_in_dim %39, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %41 = stablehlo.concatenate %38, %40, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %42 = "stablehlo.gather"(%1, %41) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %43 = stablehlo.pad %36, %42, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %44 = stablehlo.constant dense<3> : tensor<i32> loc(#loc185)
    %45 = stablehlo.broadcast_in_dim %44, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %46 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %47 = stablehlo.broadcast_in_dim %46, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %48 = stablehlo.concatenate %45, %47, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %49 = "stablehlo.gather"(%1, %48) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %50 = stablehlo.pad %43, %49, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %51 = stablehlo.constant dense<3> : tensor<i32> loc(#loc185)
    %52 = stablehlo.broadcast_in_dim %51, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %53 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %54 = stablehlo.broadcast_in_dim %53, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %55 = stablehlo.concatenate %52, %54, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %56 = "stablehlo.gather"(%1, %55) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %57 = stablehlo.pad %50, %56, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %58 = stablehlo.constant dense<4> : tensor<i32> loc(#loc185)
    %59 = stablehlo.broadcast_in_dim %58, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %60 = stablehlo.constant dense<0> : tensor<i32> loc(#loc185)
    %61 = stablehlo.broadcast_in_dim %60, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %62 = stablehlo.concatenate %59, %61, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %63 = "stablehlo.gather"(%1, %62) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %64 = stablehlo.pad %57, %63, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    %65 = stablehlo.constant dense<4> : tensor<i32> loc(#loc185)
    %66 = stablehlo.broadcast_in_dim %65, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %67 = stablehlo.constant dense<1> : tensor<i32> loc(#loc185)
    %68 = stablehlo.broadcast_in_dim %67, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5332)
    %69 = stablehlo.concatenate %66, %68, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5333)
    %70 = "stablehlo.gather"(%1, %69) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<5x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5334)
    %71 = stablehlo.pad %64, %70, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc5335)
    return %71 : tensor<1x1x1x2048x2048xbf16> loc(#loc185)
  } loc(#loc185)
  func.func private @remainder(%arg0: tensor<i32> loc(unknown)) -> (tensor<i32>, tensor<i1>) {
    %0 = stablehlo.convert %arg0 : tensor<i32> loc(#loc5336)
    %1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc188)
    %2 = stablehlo.compare  EQ, %0, %1,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5337)
    %3 = stablehlo.constant dense<1> : tensor<i32> loc(#loc188)
    %4 = call @_where(%2, %3, %0) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5338)
    %5 = stablehlo.constant dense<0> : tensor<i32> loc(#loc188)
    %6 = stablehlo.compare  LT, %4, %5,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5339)
    return %4, %6 : tensor<i32>, tensor<i1> loc(#loc188)
  } loc(#loc188)
  func.func private @_where(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc5340)
    return %0 : tensor<i32> loc(#loc5338)
  } loc(#loc5338)
  func.func private @remainder_2(%arg0: tensor<i32> loc(unknown)) -> (tensor<i32>, tensor<i1>) {
    %0 = stablehlo.convert %arg0 : tensor<i32> loc(#loc5341)
    %1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc189)
    %2 = stablehlo.compare  EQ, %0, %1,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5342)
    %3 = stablehlo.constant dense<1> : tensor<i32> loc(#loc189)
    %4 = call @_where_3(%2, %3, %0) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5343)
    %5 = stablehlo.constant dense<0> : tensor<i32> loc(#loc189)
    %6 = stablehlo.compare  LT, %4, %5,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5344)
    return %4, %6 : tensor<i32>, tensor<i1> loc(#loc189)
  } loc(#loc189)
  func.func private @_where_3(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc5345)
    return %0 : tensor<i32> loc(#loc5343)
  } loc(#loc5343)
  func.func private @softplus() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc215)
    return %0 : tensor<bf16> loc(#loc215)
  } loc(#loc215)
  func.func private @_where_4(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5346)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5347)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc217)
  } loc(#loc217)
  func.func private @softplus_5() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc220)
    return %0 : tensor<bf16> loc(#loc220)
  } loc(#loc220)
  func.func private @_where_6(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5348)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5349)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc222)
  } loc(#loc222)
  func.func private @softplus_7() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc225)
    return %0 : tensor<bf16> loc(#loc225)
  } loc(#loc225)
  func.func private @_where_8(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5350)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5351)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc227)
  } loc(#loc227)
  func.func private @softplus_9() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc230)
    return %0 : tensor<bf16> loc(#loc230)
  } loc(#loc230)
  func.func private @_where_10(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5352)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5353)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc232)
  } loc(#loc232)
  func.func private @softplus_11() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc235)
    return %0 : tensor<bf16> loc(#loc235)
  } loc(#loc235)
  func.func private @_where_12(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5354)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5355)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc237)
  } loc(#loc237)
  func.func private @softplus_13() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc240)
    return %0 : tensor<bf16> loc(#loc240)
  } loc(#loc240)
  func.func private @_where_14(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5356)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5357)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc242)
  } loc(#loc242)
  func.func private @softplus_15() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc245)
    return %0 : tensor<bf16> loc(#loc245)
  } loc(#loc245)
  func.func private @_where_16(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5358)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5359)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc247)
  } loc(#loc247)
  func.func private @softplus_17() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc250)
    return %0 : tensor<bf16> loc(#loc250)
  } loc(#loc250)
  func.func private @_where_18(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5360)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5361)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc252)
  } loc(#loc252)
  func.func private @softplus_19() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc255)
    return %0 : tensor<bf16> loc(#loc255)
  } loc(#loc255)
  func.func private @_where_20(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5362)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5363)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc257)
  } loc(#loc257)
  func.func private @softplus_21() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc260)
    return %0 : tensor<bf16> loc(#loc260)
  } loc(#loc260)
  func.func private @_where_22(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5364)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5365)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc262)
  } loc(#loc262)
  func.func private @softplus_23() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc265)
    return %0 : tensor<bf16> loc(#loc265)
  } loc(#loc265)
  func.func private @_where_24(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5366)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5367)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc267)
  } loc(#loc267)
  func.func private @softplus_25() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc270)
    return %0 : tensor<bf16> loc(#loc270)
  } loc(#loc270)
  func.func private @_where_26(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5368)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5369)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc272)
  } loc(#loc272)
  func.func private @softplus_27() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc275)
    return %0 : tensor<bf16> loc(#loc275)
  } loc(#loc275)
  func.func private @_where_28(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5370)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5371)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc277)
  } loc(#loc277)
  func.func private @softplus_29() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc280)
    return %0 : tensor<bf16> loc(#loc280)
  } loc(#loc280)
  func.func private @_where_30(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5372)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5373)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc282)
  } loc(#loc282)
  func.func private @softplus_31() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc285)
    return %0 : tensor<bf16> loc(#loc285)
  } loc(#loc285)
  func.func private @_where_32(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5374)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5375)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc287)
  } loc(#loc287)
  func.func private @softplus_33() -> tensor<bf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc290)
    return %0 : tensor<bf16> loc(#loc290)
  } loc(#loc290)
  func.func private @_where_34(%arg0: tensor<f32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5376)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5377)
    return %1 : tensor<1x1x32x2048x2048xf32> loc(#loc292)
  } loc(#loc292)
  func.func private @_where_35(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_36(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_37(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_38(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_39(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_40(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_41(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_42(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_43(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_44(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_45(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_46(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_47(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_48(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_49(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_50(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_51(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_52(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_53(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_54(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_55(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_56(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_57(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_58(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_59(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_60(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_61(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_62(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_63(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_64(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_65(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_66(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_67(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_68(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_69(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_70(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_71(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_72(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_73(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_74(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_75(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_76(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_77(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_78(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_79(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_80(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_81(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_where_82(%arg0: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc5378)
    return %0 : tensor<1xf32> loc(#loc293)
  } loc(#loc293)
  func.func private @_pad_83(%arg0: tensor<i32> loc(unknown)) -> (tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i32>) -> tensor<4x2xi32> loc(#loc5379)
    %1 = stablehlo.convert %0 : (tensor<4x2xi32>) -> tensor<4x2xbf16> loc(#loc5380)
    %2 = stablehlo.constant dense<0> : tensor<i32> loc(#loc294)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %4 = stablehlo.constant dense<0> : tensor<i32> loc(#loc294)
    %5 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %6 = stablehlo.concatenate %3, %5, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %7 = "stablehlo.gather"(%1, %6) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %8 = stablehlo.constant dense<0> : tensor<i32> loc(#loc294)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %10 = stablehlo.constant dense<1> : tensor<i32> loc(#loc294)
    %11 = stablehlo.broadcast_in_dim %10, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %12 = stablehlo.concatenate %9, %11, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %13 = "stablehlo.gather"(%1, %12) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %14 = stablehlo.constant dense<1> : tensor<i32> loc(#loc294)
    %15 = stablehlo.broadcast_in_dim %14, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %16 = stablehlo.constant dense<0> : tensor<i32> loc(#loc294)
    %17 = stablehlo.broadcast_in_dim %16, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %18 = stablehlo.concatenate %15, %17, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %19 = "stablehlo.gather"(%1, %18) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %20 = stablehlo.constant dense<1> : tensor<i32> loc(#loc294)
    %21 = stablehlo.broadcast_in_dim %20, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %22 = stablehlo.constant dense<1> : tensor<i32> loc(#loc294)
    %23 = stablehlo.broadcast_in_dim %22, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %24 = stablehlo.concatenate %21, %23, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %25 = "stablehlo.gather"(%1, %24) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %26 = stablehlo.constant dense<2> : tensor<i32> loc(#loc294)
    %27 = stablehlo.broadcast_in_dim %26, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %28 = stablehlo.constant dense<0> : tensor<i32> loc(#loc294)
    %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %30 = stablehlo.concatenate %27, %29, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %31 = "stablehlo.gather"(%1, %30) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %32 = stablehlo.constant dense<2> : tensor<i32> loc(#loc294)
    %33 = stablehlo.broadcast_in_dim %32, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %34 = stablehlo.constant dense<1> : tensor<i32> loc(#loc294)
    %35 = stablehlo.broadcast_in_dim %34, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %36 = stablehlo.concatenate %33, %35, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %37 = "stablehlo.gather"(%1, %36) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %38 = stablehlo.constant dense<3> : tensor<i32> loc(#loc294)
    %39 = stablehlo.broadcast_in_dim %38, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %40 = stablehlo.constant dense<0> : tensor<i32> loc(#loc294)
    %41 = stablehlo.broadcast_in_dim %40, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %42 = stablehlo.concatenate %39, %41, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %43 = "stablehlo.gather"(%1, %42) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    %44 = stablehlo.constant dense<3> : tensor<i32> loc(#loc294)
    %45 = stablehlo.broadcast_in_dim %44, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %46 = stablehlo.constant dense<1> : tensor<i32> loc(#loc294)
    %47 = stablehlo.broadcast_in_dim %46, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5381)
    %48 = stablehlo.concatenate %45, %47, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc5382)
    %49 = "stablehlo.gather"(%1, %48) {dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1]>, indices_are_sorted = true, slice_sizes = dense<1> : tensor<2xi64>} : (tensor<4x2xbf16>, tensor<2xi32>) -> tensor<bf16> loc(#loc5383)
    return %7, %13, %19, %25, %31, %37, %43, %49 : tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16>, tensor<bf16> loc(#loc294)
  } loc(#loc294)
  func.func private @remainder_84(%arg0: tensor<1xi32> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i1> loc(unknown)) -> tensor<1xi32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5384)
    %1 = stablehlo.remainder %arg0, %0 : tensor<1xi32> loc(#loc5384)
    %2 = stablehlo.constant dense<0> : tensor<i32> loc(#loc306)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5385)
    %4 = stablehlo.compare  LT, %1, %3,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc5385)
    %5 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc5386)
    %6 = stablehlo.compare  NE, %4, %5,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc5386)
    %7 = stablehlo.constant dense<0> : tensor<i32> loc(#loc306)
    %8 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5386)
    %9 = stablehlo.compare  NE, %1, %8,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc5386)
    %10 = stablehlo.and %6, %9 : tensor<1xi1> loc(#loc5387)
    %11 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5388)
    %12 = stablehlo.add %1, %11 : tensor<1xi32> loc(#loc5388)
    %13 = stablehlo.select %10, %12, %1 : tensor<1xi1>, tensor<1xi32> loc(#loc5389)
    return %13 : tensor<1xi32> loc(#loc306)
  } loc(#loc306)
  func.func private @remainder_85(%arg0: tensor<i32> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i1> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.remainder %arg0, %arg1 : tensor<i32> loc(#loc5390)
    %1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc350)
    %2 = stablehlo.compare  LT, %0, %1,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5391)
    %3 = stablehlo.compare  NE, %2, %arg2,  UNSIGNED : (tensor<i1>, tensor<i1>) -> tensor<i1> loc(#loc5392)
    %4 = stablehlo.constant dense<0> : tensor<i32> loc(#loc350)
    %5 = stablehlo.compare  NE, %0, %4,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5392)
    %6 = stablehlo.and %3, %5 : tensor<i1> loc(#loc5393)
    %7 = stablehlo.add %0, %arg1 : tensor<i32> loc(#loc5394)
    %8 = stablehlo.select %6, %7, %0 : tensor<i1>, tensor<i32> loc(#loc5395)
    return %8 : tensor<i32> loc(#loc350)
  } loc(#loc350)
  func.func private @_where_86(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xi1> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.select %arg2, %arg0, %arg1 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xbf16> loc(#loc5396)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc356)
  } loc(#loc356)
  func.func private @_where_87(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_88(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_89(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_90(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5400)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5401)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc378)
  } loc(#loc378)
  func.func private @_where_91(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_92(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5403)
    return %0 : tensor<1x64xbf16> loc(#loc382)
  } loc(#loc382)
  func.func private @logaddexp(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5404)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5404)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5405)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5406)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5406)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5407)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5408)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5409)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5410)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5411)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5411)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5411)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5412)
    return %12 : tensor<1x64xbf16> loc(#loc5403)
  } loc(#loc5403)
  func.func private @_einsum_93(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5413)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc387)
  } loc(#loc387)
  func.func private @_where_94(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5414)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5415)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5416)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc392)
  } loc(#loc392)
  func.func private @_einsum_95(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5417)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5418)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc401)
  } loc(#loc401)
  func.func private @_where_96(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_97(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5420)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc402)
  } loc(#loc402)
  func.func private @_where_98(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_99(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_100(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_101(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5422)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc420)
  } loc(#loc420)
  func.func private @_where_102(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_103(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_104(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5425)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc435)
  } loc(#loc435)
  func.func private @_where_105(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_106(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_107(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_108(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_109(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5426)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5427)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc458)
  } loc(#loc458)
  func.func private @_where_110(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_111(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_112(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5428)
    return %0 : tensor<1x64xbf16> loc(#loc462)
  } loc(#loc462)
  func.func private @logaddexp_112(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5429)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5429)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5430)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5431)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5431)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5432)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5433)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5434)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5435)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5436)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5436)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5436)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5437)
    return %12 : tensor<1x64xbf16> loc(#loc5428)
  } loc(#loc5428)
  func.func private @_einsum_113(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5438)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc467)
  } loc(#loc467)
  func.func private @_where_114(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5439)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5440)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5441)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc472)
  } loc(#loc472)
  func.func private @_einsum_115(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5442)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5443)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc481)
  } loc(#loc481)
  func.func private @_where_116(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_117(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5444)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc482)
  } loc(#loc482)
  func.func private @_where_118(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_119(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_120(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_121(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5445)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc500)
  } loc(#loc500)
  func.func private @_where_122(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_123(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_124(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5446)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc512)
  } loc(#loc512)
  func.func private @_where_125(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_126(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_127(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_128(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_129(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5447)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5448)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc535)
  } loc(#loc535)
  func.func private @_where_130(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_131(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_132(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5449)
    return %0 : tensor<1x64xbf16> loc(#loc539)
  } loc(#loc539)
  func.func private @logaddexp_132(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5450)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5450)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5451)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5452)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5452)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5453)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5454)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5455)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5456)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5457)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5457)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5457)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5458)
    return %12 : tensor<1x64xbf16> loc(#loc5449)
  } loc(#loc5449)
  func.func private @_einsum_133(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5459)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc544)
  } loc(#loc544)
  func.func private @_where_134(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5460)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5461)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5462)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc549)
  } loc(#loc549)
  func.func private @_einsum_135(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5463)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5464)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc558)
  } loc(#loc558)
  func.func private @_where_136(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_137(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5465)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc559)
  } loc(#loc559)
  func.func private @_where_138(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_139(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_140(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_141(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5466)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc577)
  } loc(#loc577)
  func.func private @_where_142(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_143(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_144(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5467)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc589)
  } loc(#loc589)
  func.func private @_where_145(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_146(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_147(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_148(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_149(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5468)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5469)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc612)
  } loc(#loc612)
  func.func private @_where_150(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_151(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_152(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5470)
    return %0 : tensor<1x64xbf16> loc(#loc616)
  } loc(#loc616)
  func.func private @logaddexp_152(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5471)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5471)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5472)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5473)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5473)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5474)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5475)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5476)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5477)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5478)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5478)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5478)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5479)
    return %12 : tensor<1x64xbf16> loc(#loc5470)
  } loc(#loc5470)
  func.func private @_einsum_153(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5480)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc621)
  } loc(#loc621)
  func.func private @_where_154(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5481)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5482)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5483)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc626)
  } loc(#loc626)
  func.func private @_einsum_155(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5484)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5485)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc635)
  } loc(#loc635)
  func.func private @_where_156(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_157(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5486)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc636)
  } loc(#loc636)
  func.func private @_where_158(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_159(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_160(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_161(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5487)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc654)
  } loc(#loc654)
  func.func private @_where_162(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_163(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_164(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5488)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc666)
  } loc(#loc666)
  func.func private @_where_165(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_166(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_167(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_168(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_169(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5489)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5490)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc689)
  } loc(#loc689)
  func.func private @_where_170(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_171(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_172(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5491)
    return %0 : tensor<1x64xbf16> loc(#loc693)
  } loc(#loc693)
  func.func private @logaddexp_172(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5492)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5492)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5493)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5494)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5494)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5495)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5496)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5497)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5498)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5499)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5499)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5499)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5500)
    return %12 : tensor<1x64xbf16> loc(#loc5491)
  } loc(#loc5491)
  func.func private @_einsum_173(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5501)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc698)
  } loc(#loc698)
  func.func private @_where_174(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5502)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5503)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5504)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc703)
  } loc(#loc703)
  func.func private @_einsum_175(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5505)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5506)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc712)
  } loc(#loc712)
  func.func private @_where_176(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_177(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5507)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc713)
  } loc(#loc713)
  func.func private @_where_178(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_179(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_180(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_181(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5508)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc731)
  } loc(#loc731)
  func.func private @_where_182(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_183(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_184(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5509)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc743)
  } loc(#loc743)
  func.func private @_where_185(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_186(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_187(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_188(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_189(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5510)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5511)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc766)
  } loc(#loc766)
  func.func private @_where_190(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_191(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_192(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5512)
    return %0 : tensor<1x64xbf16> loc(#loc770)
  } loc(#loc770)
  func.func private @logaddexp_192(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5513)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5513)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5514)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5515)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5515)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5516)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5517)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5518)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5519)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5520)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5520)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5520)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5521)
    return %12 : tensor<1x64xbf16> loc(#loc5512)
  } loc(#loc5512)
  func.func private @_einsum_193(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5522)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc775)
  } loc(#loc775)
  func.func private @_where_194(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5523)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5524)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5525)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc780)
  } loc(#loc780)
  func.func private @_einsum_195(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5526)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5527)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc789)
  } loc(#loc789)
  func.func private @_where_196(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_197(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5528)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc790)
  } loc(#loc790)
  func.func private @_where_198(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_199(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_200(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_201(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5529)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc808)
  } loc(#loc808)
  func.func private @_where_202(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_203(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_204(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5530)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc820)
  } loc(#loc820)
  func.func private @_where_205(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_206(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_207(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_208(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_209(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5531)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5532)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc843)
  } loc(#loc843)
  func.func private @_where_210(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_211(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_212(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5533)
    return %0 : tensor<1x64xbf16> loc(#loc847)
  } loc(#loc847)
  func.func private @logaddexp_212(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5534)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5534)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5535)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5536)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5536)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5537)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5538)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5539)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5540)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5541)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5541)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5541)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5542)
    return %12 : tensor<1x64xbf16> loc(#loc5533)
  } loc(#loc5533)
  func.func private @_einsum_213(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5543)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc852)
  } loc(#loc852)
  func.func private @_where_214(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5544)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5545)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5546)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc857)
  } loc(#loc857)
  func.func private @_einsum_215(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5547)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5548)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc866)
  } loc(#loc866)
  func.func private @_where_216(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_217(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5549)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc867)
  } loc(#loc867)
  func.func private @_where_218(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_219(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_220(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_221(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5550)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc885)
  } loc(#loc885)
  func.func private @_where_222(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_223(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_224(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5551)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc897)
  } loc(#loc897)
  func.func private @_where_225(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_226(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_227(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_228(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_229(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5552)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5553)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc920)
  } loc(#loc920)
  func.func private @_where_230(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_231(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_232(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5554)
    return %0 : tensor<1x64xbf16> loc(#loc924)
  } loc(#loc924)
  func.func private @logaddexp_232(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5555)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5555)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5556)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5557)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5557)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5558)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5559)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5560)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5561)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5562)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5562)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5562)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5563)
    return %12 : tensor<1x64xbf16> loc(#loc5554)
  } loc(#loc5554)
  func.func private @_einsum_233(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5564)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc929)
  } loc(#loc929)
  func.func private @_where_234(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5565)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5566)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5567)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc934)
  } loc(#loc934)
  func.func private @_einsum_235(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5568)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5569)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc943)
  } loc(#loc943)
  func.func private @_where_236(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_237(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5570)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc944)
  } loc(#loc944)
  func.func private @_where_238(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_239(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_240(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_241(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5571)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc962)
  } loc(#loc962)
  func.func private @_where_242(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_243(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_244(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5572)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc974)
  } loc(#loc974)
  func.func private @_where_245(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_246(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_247(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_248(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_249(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5573)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5574)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc997)
  } loc(#loc997)
  func.func private @_where_250(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_251(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_252(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5575)
    return %0 : tensor<1x64xbf16> loc(#loc1001)
  } loc(#loc1001)
  func.func private @logaddexp_252(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5576)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5576)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5577)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5578)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5578)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5579)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5580)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5581)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5582)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5583)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5583)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5583)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5584)
    return %12 : tensor<1x64xbf16> loc(#loc5575)
  } loc(#loc5575)
  func.func private @_einsum_253(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5585)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1006)
  } loc(#loc1006)
  func.func private @_where_254(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5586)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5587)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5588)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1011)
  } loc(#loc1011)
  func.func private @_einsum_255(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5589)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5590)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1020)
  } loc(#loc1020)
  func.func private @_where_256(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_257(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5591)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1021)
  } loc(#loc1021)
  func.func private @_where_258(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_259(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_260(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_261(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5592)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1039)
  } loc(#loc1039)
  func.func private @_where_262(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_263(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_264(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5593)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1051)
  } loc(#loc1051)
  func.func private @_where_265(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_266(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_267(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_268(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_269(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5594)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5595)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1074)
  } loc(#loc1074)
  func.func private @_where_270(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_271(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_272(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5596)
    return %0 : tensor<1x64xbf16> loc(#loc1078)
  } loc(#loc1078)
  func.func private @logaddexp_272(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5597)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5597)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5598)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5599)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5599)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5600)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5601)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5602)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5603)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5604)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5604)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5604)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5605)
    return %12 : tensor<1x64xbf16> loc(#loc5596)
  } loc(#loc5596)
  func.func private @_einsum_273(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5606)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1083)
  } loc(#loc1083)
  func.func private @_where_274(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5607)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5608)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5609)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1088)
  } loc(#loc1088)
  func.func private @_einsum_275(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5610)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5611)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1097)
  } loc(#loc1097)
  func.func private @_where_276(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_277(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5612)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1098)
  } loc(#loc1098)
  func.func private @_where_278(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_279(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_280(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_281(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5613)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1116)
  } loc(#loc1116)
  func.func private @_where_282(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_283(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_284(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5614)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1128)
  } loc(#loc1128)
  func.func private @_where_285(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_286(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_287(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_288(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_289(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5615)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5616)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1151)
  } loc(#loc1151)
  func.func private @_where_290(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_291(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_292(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5617)
    return %0 : tensor<1x64xbf16> loc(#loc1155)
  } loc(#loc1155)
  func.func private @logaddexp_292(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5618)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5618)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5619)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5620)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5620)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5621)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5622)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5623)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5624)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5625)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5625)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5625)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5626)
    return %12 : tensor<1x64xbf16> loc(#loc5617)
  } loc(#loc5617)
  func.func private @_einsum_293(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5627)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1160)
  } loc(#loc1160)
  func.func private @_where_294(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5628)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5629)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5630)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1165)
  } loc(#loc1165)
  func.func private @_einsum_295(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5631)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5632)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1174)
  } loc(#loc1174)
  func.func private @_where_296(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_297(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5633)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1175)
  } loc(#loc1175)
  func.func private @_where_298(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_299(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_300(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_301(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5634)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1193)
  } loc(#loc1193)
  func.func private @_where_302(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_303(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_304(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5635)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1205)
  } loc(#loc1205)
  func.func private @_where_305(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_306(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_307(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_308(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_309(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5636)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5637)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1228)
  } loc(#loc1228)
  func.func private @_where_310(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_311(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_312(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5638)
    return %0 : tensor<1x64xbf16> loc(#loc1232)
  } loc(#loc1232)
  func.func private @logaddexp_312(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5639)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5639)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5640)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5641)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5641)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5642)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5643)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5644)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5645)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5646)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5646)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5646)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5647)
    return %12 : tensor<1x64xbf16> loc(#loc5638)
  } loc(#loc5638)
  func.func private @_einsum_313(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5648)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1237)
  } loc(#loc1237)
  func.func private @_where_314(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5649)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5650)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5651)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1242)
  } loc(#loc1242)
  func.func private @_einsum_315(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5652)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5653)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1251)
  } loc(#loc1251)
  func.func private @_where_316(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_317(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5654)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1252)
  } loc(#loc1252)
  func.func private @_where_318(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_319(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_320(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_321(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5655)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1270)
  } loc(#loc1270)
  func.func private @_where_322(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_323(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_324(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5656)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1282)
  } loc(#loc1282)
  func.func private @_where_325(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_326(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_327(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_328(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_329(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5657)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5658)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1305)
  } loc(#loc1305)
  func.func private @_where_330(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_331(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_332(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5659)
    return %0 : tensor<1x64xbf16> loc(#loc1309)
  } loc(#loc1309)
  func.func private @logaddexp_332(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5660)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5660)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5661)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5662)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5662)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5663)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5664)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5665)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5666)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5667)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5667)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5667)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5668)
    return %12 : tensor<1x64xbf16> loc(#loc5659)
  } loc(#loc5659)
  func.func private @_einsum_333(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5669)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1314)
  } loc(#loc1314)
  func.func private @_where_334(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5670)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5671)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5672)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1319)
  } loc(#loc1319)
  func.func private @_einsum_335(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5673)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5674)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1328)
  } loc(#loc1328)
  func.func private @_where_336(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_337(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5675)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1329)
  } loc(#loc1329)
  func.func private @_where_338(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_339(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_340(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_341(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5676)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1347)
  } loc(#loc1347)
  func.func private @_where_342(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_343(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_344(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5677)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1359)
  } loc(#loc1359)
  func.func private @_where_345(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_346(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_347(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_348(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_349(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5678)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5679)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1382)
  } loc(#loc1382)
  func.func private @_where_350(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_351(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_352(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5680)
    return %0 : tensor<1x64xbf16> loc(#loc1386)
  } loc(#loc1386)
  func.func private @logaddexp_352(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5681)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5681)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5682)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5683)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5683)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5684)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5685)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5686)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5687)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5688)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5688)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5688)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5689)
    return %12 : tensor<1x64xbf16> loc(#loc5680)
  } loc(#loc5680)
  func.func private @_einsum_353(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5690)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1391)
  } loc(#loc1391)
  func.func private @_where_354(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5691)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5692)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5693)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1396)
  } loc(#loc1396)
  func.func private @_einsum_355(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5694)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5695)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1405)
  } loc(#loc1405)
  func.func private @_where_356(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_357(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5696)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1406)
  } loc(#loc1406)
  func.func private @_where_358(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_359(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_360(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_361(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5697)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1424)
  } loc(#loc1424)
  func.func private @_where_362(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_363(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_364(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5698)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1436)
  } loc(#loc1436)
  func.func private @_where_365(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_366(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_367(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_368(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_369(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5699)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5700)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1459)
  } loc(#loc1459)
  func.func private @_where_370(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_371(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_372(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5701)
    return %0 : tensor<1x64xbf16> loc(#loc1463)
  } loc(#loc1463)
  func.func private @logaddexp_372(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5702)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5702)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5703)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5704)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5704)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5705)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5706)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5707)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5708)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5709)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5709)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5709)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5710)
    return %12 : tensor<1x64xbf16> loc(#loc5701)
  } loc(#loc5701)
  func.func private @_einsum_373(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5711)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1468)
  } loc(#loc1468)
  func.func private @_where_374(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5712)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5713)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5714)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1473)
  } loc(#loc1473)
  func.func private @_einsum_375(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5715)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5716)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1482)
  } loc(#loc1482)
  func.func private @_where_376(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_377(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5717)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1483)
  } loc(#loc1483)
  func.func private @_where_378(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_379(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_380(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_381(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5718)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1501)
  } loc(#loc1501)
  func.func private @_where_382(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_383(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_384(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5719)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1513)
  } loc(#loc1513)
  func.func private @_where_385(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_386(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_387(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_388(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5399)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x3x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_389(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5720)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5721)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc1536)
  } loc(#loc1536)
  func.func private @_where_390(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5402)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5398)
    return %1 : tensor<1x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @softplus_391(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = call @logaddexp_392(%arg0, %arg1) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5722)
    return %0 : tensor<1x64xbf16> loc(#loc1540)
  } loc(#loc1540)
  func.func private @logaddexp_392(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5723)
    %1 = stablehlo.subtract %arg0, %0 : tensor<1x64xbf16> loc(#loc5723)
    %2 = stablehlo.compare  NE, %1, %1,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5724)
    %3 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5725)
    %4 = stablehlo.maximum %arg0, %3 : tensor<1x64xbf16> loc(#loc5725)
    %5 = stablehlo.abs %1 : tensor<1x64xbf16> loc(#loc5726)
    %6 = stablehlo.negate %5 : tensor<1x64xbf16> loc(#loc5727)
    %7 = stablehlo.exponential %6 : tensor<1x64xbf16> loc(#loc5728)
    %8 = stablehlo.log_plus_one %7 : tensor<1x64xbf16> loc(#loc5729)
    %9 = stablehlo.add %4, %8 : tensor<1x64xbf16> loc(#loc5730)
    %10 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5730)
    %11 = stablehlo.add %arg0, %10 : tensor<1x64xbf16> loc(#loc5730)
    %12 = stablehlo.select %2, %11, %9 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5731)
    return %12 : tensor<1x64xbf16> loc(#loc5722)
  } loc(#loc5722)
  func.func private @_einsum_393(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5732)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc1545)
  } loc(#loc1545)
  func.func private @_where_394(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5733)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5734)
    %2 = stablehlo.select %1, %arg1, %arg2 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5735)
    return %2 : tensor<1x1x32x2048x2048xf32> loc(#loc1550)
  } loc(#loc1550)
  func.func private @_einsum_395(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5736)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5737)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc1559)
  } loc(#loc1559)
  func.func private @_where_396(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5419)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x32x64xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_397(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5738)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1560)
  } loc(#loc1560)
  func.func private @_where_398(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_399(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_400(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5421)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x2048x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_401(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5739)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc1578)
  } loc(#loc1578)
  func.func private @_where_402(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5423)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5398)
    return %1 : tensor<1x8192xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_where_403(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5424)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x8192x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_einsum_404(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5740)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1590)
  } loc(#loc1590)
  func.func private @_where_405(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5397)
    %1 = stablehlo.select %0, %arg2, %arg1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5398)
    return %1 : tensor<1x2048xbf16> loc(#loc372)
  } loc(#loc372)
  func.func private @_pad_406(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown), %arg2: tensor<bf16> loc(unknown), %arg3: tensor<bf16> loc(unknown), %arg4: tensor<bf16> loc(unknown), %arg5: tensor<bf16> loc(unknown), %arg6: tensor<bf16> loc(unknown), %arg7: tensor<bf16> loc(unknown), %arg8: tensor<bf16> loc(unknown)) -> tensor<2x1x2048x2048xbf16> {
    %0 = stablehlo.pad %arg0, %arg1, low = [1, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5741)
    %1 = stablehlo.pad %0, %arg2, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    %2 = stablehlo.pad %1, %arg3, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    %3 = stablehlo.pad %2, %arg4, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    %4 = stablehlo.pad %3, %arg5, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    %5 = stablehlo.pad %4, %arg6, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    %6 = stablehlo.pad %5, %arg7, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    %7 = stablehlo.pad %6, %arg8, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc5742)
    return %7 : tensor<2x1x2048x2048xbf16> loc(#loc1598)
  } loc(#loc1598)
  func.func private @_where_407(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_408(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_409(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_410(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_411(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_412(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_413(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_414(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_415(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_416(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_417(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_418(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_419(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_420(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_421(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_422(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_423(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_424(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_425(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_426(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_427(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_428(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_429(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_430(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_431(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_432(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_433(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_434(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_435(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_436(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_437(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_438(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_439(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_440(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_441(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_442(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_443(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_444(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_445(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_446(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_447(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_448(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_449(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_450(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_451(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_452(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_453(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_where_454(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc5743)
    return %0 : tensor<1xf32> loc(#loc1609)
  } loc(#loc1609)
  func.func private @_einsum_455(%arg0: tensor<1x2048x2048xbf16> loc(unknown), %arg1: tensor<2048x51200xbf16> loc(unknown)) -> tensor<1x2048x51200xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x2048xbf16>, tensor<2048x51200xbf16>) -> tensor<1x2048x51200xbf16> loc(#loc5744)
    return %0 : tensor<1x2048x51200xbf16> loc(#loc1877)
  } loc(#loc1877)
  func.func private @log_softmax(%arg0: tensor<1x2048x51200xf32> loc(unknown)) -> (tensor<1x2048x51200xf32>, tensor<1x2048x51200xf32>, tensor<1x2048x1xf32>) {
    %0 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5745)
    %1 = stablehlo.reduce(%arg0 init: %0) across dimensions = [2] : (tensor<1x2048x51200xf32>, tensor<f32>) -> tensor<1x2048xf32>
     reducer(%arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown))  {
      %12 = stablehlo.maximum %arg1, %arg2 : tensor<f32> loc(#loc5745)
      stablehlo.return %12 : tensor<f32> loc(#loc5745)
    } loc(#loc5745)
    %2 = stablehlo.broadcast_in_dim %1, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x1xf32> loc(#loc5746)
    %3 = stablehlo.broadcast_in_dim %2, dims = [0, 1, 2] : (tensor<1x2048x1xf32>) -> tensor<1x2048x51200xf32> loc(#loc5747)
    %4 = stablehlo.subtract %arg0, %3 : tensor<1x2048x51200xf32> loc(#loc5747)
    %5 = stablehlo.exponential %4 : tensor<1x2048x51200xf32> loc(#loc5748)
    %6 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5749)
    %7 = stablehlo.reduce(%5 init: %6) across dimensions = [2] : (tensor<1x2048x51200xf32>, tensor<f32>) -> tensor<1x2048xf32>
     reducer(%arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown))  {
      %12 = stablehlo.add %arg1, %arg2 : tensor<f32> loc(#loc5749)
      stablehlo.return %12 : tensor<f32> loc(#loc5749)
    } loc(#loc5749)
    %8 = stablehlo.broadcast_in_dim %7, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x1xf32> loc(#loc5746)
    %9 = stablehlo.log %8 : tensor<1x2048x1xf32> loc(#loc5750)
    %10 = stablehlo.broadcast_in_dim %9, dims = [0, 1, 2] : (tensor<1x2048x1xf32>) -> tensor<1x2048x51200xf32> loc(#loc5747)
    %11 = stablehlo.subtract %4, %10 : tensor<1x2048x51200xf32> loc(#loc5747)
    return %11, %5, %8 : tensor<1x2048x51200xf32>, tensor<1x2048x51200xf32>, tensor<1x2048x1xf32> loc(#loc1886)
  } loc(#loc1886)
  func.func private @_one_hot_456(%arg0: tensor<1x2048xi32> loc(unknown)) -> tensor<1x2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<1x2048xi32>) -> tensor<1x2048x1xi32> loc(#loc5751)
    %1 = stablehlo.iota dim = 2 : tensor<1x1x51200xi32> loc(#loc5752)
    %2 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 2] : (tensor<1x2048x1xi32>) -> tensor<1x2048x51200xi32> loc(#loc5753)
    %3 = stablehlo.broadcast_in_dim %1, dims = [0, 1, 2] : (tensor<1x1x51200xi32>) -> tensor<1x2048x51200xi32> loc(#loc5753)
    %4 = stablehlo.compare  EQ, %2, %3,  SIGNED : (tensor<1x2048x51200xi32>, tensor<1x2048x51200xi32>) -> tensor<1x2048x51200xi1> loc(#loc5753)
    %5 = stablehlo.convert %4 : (tensor<1x2048x51200xi1>) -> tensor<1x2048x51200xf32> loc(#loc5754)
    return %5 : tensor<1x2048x51200xf32> loc(#loc1888)
  } loc(#loc1888)
  func.func private @argmax(%arg0: tensor<1x2048x51200xf32> loc(unknown)) -> tensor<1x2048xi32> {
    %0 = stablehlo.iota dim = 2 : tensor<1x2048x51200xi32> loc(#loc5755)
    %1 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc1892)
    %2 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1892)
    %3:2 = stablehlo.reduce(%arg0 init: %1), (%0 init: %2) across dimensions = [2] : (tensor<1x2048x51200xf32>, tensor<1x2048x51200xi32>, tensor<f32>, tensor<i32>) -> (tensor<1x2048xf32>, tensor<1x2048xi32>)
     reducer(%arg1: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) (%arg2: tensor<i32> loc(unknown), %arg4: tensor<i32> loc(unknown))  {
      %4 = stablehlo.compare  GT, %arg1, %arg3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5757)
      %5 = stablehlo.compare  NE, %arg1, %arg1,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5758)
      %6 = stablehlo.or %4, %5 : tensor<i1> loc(#loc5759)
      %7 = stablehlo.compare  EQ, %arg1, %arg3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc5760)
      %8 = stablehlo.compare  LT, %arg2, %arg4,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5761)
      %9 = stablehlo.and %7, %8 : tensor<i1> loc(#loc5762)
      %10 = stablehlo.or %6, %9 : tensor<i1> loc(#loc5759)
      %11 = stablehlo.select %6, %arg1, %arg3 : tensor<i1>, tensor<f32> loc(#loc5763)
      %12 = stablehlo.select %10, %arg2, %arg4 : tensor<i1>, tensor<i32> loc(#loc5763)
      stablehlo.return %11, %12 : tensor<f32>, tensor<i32> loc(#loc5756)
    } loc(#loc5756)
    return %3#1 : tensor<1x2048xi32> loc(#loc1892)
  } loc(#loc1892)
  func.func private @log_softmax_457(%arg0: tensor<1x2048x51200xf32> loc(unknown), %arg1: tensor<1x2048x1xf32> loc(unknown), %arg2: tensor<1x2048x51200xf32> loc(unknown)) -> tensor<1x2048x51200xf32> {
    %0 = stablehlo.negate %arg2 : tensor<1x2048x51200xf32> loc(#loc5764)
    %1 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5765)
    %2 = stablehlo.reduce(%0 init: %1) across dimensions = [2] : (tensor<1x2048x51200xf32>, tensor<f32>) -> tensor<1x2048xf32>
     reducer(%arg3: tensor<f32> loc(unknown), %arg4: tensor<f32> loc(unknown))  {
      %11 = stablehlo.add %arg3, %arg4 : tensor<f32> loc(#loc5765)
      stablehlo.return %11 : tensor<f32> loc(#loc5765)
    } loc(#loc5765)
    %3 = stablehlo.reshape %2 : (tensor<1x2048xf32>) -> tensor<1x2048x1xf32> loc(#loc5766)
    %4 = stablehlo.divide %3, %arg1 : tensor<1x2048x1xf32> loc(#loc5767)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5768)
    %6 = stablehlo.reduce(%4 init: %5) across dimensions = [0, 2] : (tensor<1x2048x1xf32>, tensor<f32>) -> tensor<2048xf32>
     reducer(%arg3: tensor<f32> loc(unknown), %arg4: tensor<f32> loc(unknown))  {
      %11 = stablehlo.add %arg3, %arg4 : tensor<f32> loc(#loc5768)
      stablehlo.return %11 : tensor<f32> loc(#loc5768)
    } loc(#loc5768)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc5769)
    %8 = stablehlo.broadcast_in_dim %7, dims = [0, 1] : (tensor<1x2048xf32>) -> tensor<1x2048x51200xf32> loc(#loc5770)
    %9 = stablehlo.multiply %8, %arg0 : tensor<1x2048x51200xf32> loc(#loc5771)
    %10 = stablehlo.add %arg2, %9 : tensor<1x2048x51200xf32> loc(#loc5772)
    return %10 : tensor<1x2048x51200xf32> loc(#loc1921)
  } loc(#loc1921)
  func.func private @_einsum_458(%arg0: tensor<2048x51200xbf16> loc(unknown), %arg1: tensor<1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x2048x51200xbf16> loc(unknown)) -> (tensor<1x2048x2048xbf16>, tensor<2048x51200xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, contracting_dims = [0, 1] x [0, 1], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x51200xbf16>, tensor<1x2048x2048xbf16>) -> tensor<51200x2048xbf16> loc(#loc5773)
    %1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<51200x2048xbf16>) -> tensor<2048x51200xbf16> loc(#loc5774)
    %2 = stablehlo.dot_general %arg2, %arg0, contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x51200xbf16>, tensor<2048x51200xbf16>) -> tensor<1x2048x2048xbf16> loc(#loc5775)
    return %2, %1 : tensor<1x2048x2048xbf16>, tensor<2048x51200xbf16> loc(#loc1931)
  } loc(#loc1931)
  func.func private @remainder_459(%arg0: tensor<1xi32> loc(unknown), %arg1: tensor<i32> loc(unknown)) -> tensor<1xi32> {
    %0 = stablehlo.convert %arg1 : tensor<i32> loc(#loc5776)
    %1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1986)
    %2 = stablehlo.compare  EQ, %0, %1,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5777)
    %3 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1986)
    %4 = call @_where_460(%2, %3, %0) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5778)
    %5 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5779)
    %6 = stablehlo.remainder %arg0, %5 : tensor<1xi32> loc(#loc5779)
    %7 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1986)
    %8 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5780)
    %9 = stablehlo.compare  NE, %6, %8,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc5780)
    %10 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1986)
    %11 = stablehlo.broadcast_in_dim %10, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5781)
    %12 = stablehlo.compare  LT, %6, %11,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1> loc(#loc5781)
    %13 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1986)
    %14 = stablehlo.compare  LT, %4, %13,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5781)
    %15 = stablehlo.broadcast_in_dim %14, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc5780)
    %16 = stablehlo.compare  NE, %12, %15,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc5780)
    %17 = stablehlo.and %16, %9 : tensor<1xi1> loc(#loc5782)
    %18 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<i32>) -> tensor<1xi32> loc(#loc5783)
    %19 = stablehlo.add %6, %18 : tensor<1xi32> loc(#loc5783)
    %20 = stablehlo.select %17, %19, %6 : tensor<1xi1>, tensor<1xi32> loc(#loc5784)
    return %20 : tensor<1xi32> loc(#loc1986)
  } loc(#loc1986)
  func.func private @_where_460(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc5785)
    return %0 : tensor<i32> loc(#loc5778)
  } loc(#loc5778)
  func.func private @remainder_461(%arg0: tensor<i32> loc(unknown), %arg1: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.convert %arg1 : tensor<i32> loc(#loc5786)
    %1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1987)
    %2 = stablehlo.compare  EQ, %0, %1,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5787)
    %3 = stablehlo.constant dense<1> : tensor<i32> loc(#loc1987)
    %4 = call @_where_462(%2, %3, %0) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> loc(#loc5788)
    %5 = stablehlo.remainder %arg0, %4 : tensor<i32> loc(#loc5789)
    %6 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1987)
    %7 = stablehlo.compare  NE, %5, %6,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5790)
    %8 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1987)
    %9 = stablehlo.compare  LT, %5, %8,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5791)
    %10 = stablehlo.constant dense<0> : tensor<i32> loc(#loc1987)
    %11 = stablehlo.compare  LT, %4, %10,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc5791)
    %12 = stablehlo.compare  NE, %9, %11,  UNSIGNED : (tensor<i1>, tensor<i1>) -> tensor<i1> loc(#loc5790)
    %13 = stablehlo.and %12, %7 : tensor<i1> loc(#loc5792)
    %14 = stablehlo.add %5, %4 : tensor<i32> loc(#loc5793)
    %15 = stablehlo.select %13, %14, %5 : tensor<i1>, tensor<i32> loc(#loc5794)
    return %15 : tensor<i32> loc(#loc1987)
  } loc(#loc1987)
  func.func private @_where_462(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc5795)
    return %0 : tensor<i32> loc(#loc5788)
  } loc(#loc5788)
  func.func private @_where_463(%arg0: tensor<1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xbf16> loc(#loc5796)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc1995)
  } loc(#loc1995)
  func.func private @_where_464(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_465(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_466(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_467(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_468(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_469(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_470(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_471(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_472(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_473(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_474(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_475(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_476(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_477(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_478(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_479(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_480(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_481(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_482(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_483(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_484(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_485(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_486(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_487(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_488(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_489(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_490(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_491(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_492(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_493(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_494(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_495(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_496(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_497(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_498(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_499(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_500(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_501(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_502(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_503(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_504(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_505(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_506(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_507(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_508(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_509(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_510(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_511(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_512(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_513(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_514(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_515(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_516(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_517(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_518(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_519(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_520(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_521(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_522(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_523(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_524(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_525(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_526(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_527(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_528(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_529(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_530(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_531(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_532(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_533(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_534(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_535(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_536(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_537(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_538(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_539(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_540(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_541(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_542(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_543(%arg0: tensor<1xi1> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    return %0 : tensor<1x2048xi1> loc(#loc2017)
  } loc(#loc2017)
  func.func private @_where_544(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_545(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_546(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_547(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_548(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_549(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_550(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_551(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_552(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_553(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_554(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_555(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_556(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_557(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_558(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_559(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_560(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_561(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_562(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_563(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_564(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_565(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_566(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_567(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_568(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_569(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_570(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_571(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_572(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_573(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_574(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_575(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_576(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_577(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_578(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_579(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_580(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_581(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_582(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_583(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_584(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_585(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_586(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_587(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_588(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_589(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_590(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_591(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_592(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_593(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_594(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_595(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_596(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_597(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_598(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_599(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_600(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_601(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_602(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_603(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_604(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_605(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_606(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_607(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_608(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_609(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_610(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_611(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_612(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_613(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_614(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_615(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_616(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_617(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_618(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_619(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_620(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_621(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_622(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_623(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_624(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_625(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_626(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_627(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_628(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_629(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown), %arg2: tensor<1x8192xbf16> loc(unknown)) -> (tensor<1x8192xbf16>, tensor<1x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192xi1> loc(#loc5797)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192xbf16>, tensor<1x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_630(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x2048x8192xbf16> loc(unknown)) -> (tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x8192xi1> loc(#loc5799)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x8192xbf16>, tensor<1x2048x8192xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_631(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_632(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown), %arg2: tensor<1x8192x2048xbf16> loc(unknown)) -> (tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x8192x2048xi1> loc(#loc5801)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x8192x2048xbf16>, tensor<1x8192x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_633(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_634(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_635(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_636(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown), %arg2: tensor<1x2048xbf16> loc(unknown)) -> (tensor<1x2048xbf16>, tensor<1x2048xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048xi1> loc(#loc5800)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048xbf16>, tensor<1x2048xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_637(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> (tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc5802)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x3x2048x32x64xbf16>, tensor<1x3x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_638(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown), %arg2: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x64xi1> loc(#loc5803)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x64xbf16>, tensor<1x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_where_639(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1>) {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<1xi1>) -> tensor<1x2048x32x64xi1> loc(#loc5804)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %1, %0 : tensor<1x2048x32x64xbf16>, tensor<1x2048x32x64xi1> loc(#loc2016)
  } loc(#loc2016)
  func.func private @_einsum_640(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5805)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5806)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2075)
  } loc(#loc2075)
  func.func private @softplus_641(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2081)
    %1 = call @logaddexp_642(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5807)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2081)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5808)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5808)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2081)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5809)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5810)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5811)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2081)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5808)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5808)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2081)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5809)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5810)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5811)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5812)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5813)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2081)
  } loc(#loc2081)
  func.func private @logaddexp_642(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5814)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5814)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5815)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5815)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5816)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5817)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5817)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5818)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5819)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5820)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5821)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5817)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5822)
    return %12 : tensor<1x64xbf16> loc(#loc5807)
  } loc(#loc5807)
  func.func private @_einsum_643(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5823)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2085)
  } loc(#loc2085)
  func.func private @_where_644(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5824)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5825)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5826)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5827)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5828)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2094)
  } loc(#loc2094)
  func.func private @_einsum_645(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5829)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5830)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2105)
  } loc(#loc2105)
  func.func private @_einsum_646(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5831)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2106)
  } loc(#loc2106)
  func.func private @_einsum_647(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5832)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2126)
  } loc(#loc2126)
  func.func private @_einsum_648(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5833)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2138)
  } loc(#loc2138)
  func.func private @_einsum_649(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5834)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5835)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2185)
  } loc(#loc2185)
  func.func private @softplus_650(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2191)
    %1 = call @logaddexp_651(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5836)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2191)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5837)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5837)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2191)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5838)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5839)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5840)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2191)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5837)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5837)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2191)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5838)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5839)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5840)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5841)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5842)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2191)
  } loc(#loc2191)
  func.func private @logaddexp_651(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5843)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5843)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5844)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5844)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5845)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5846)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5846)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5847)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5848)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5849)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5850)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5846)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5851)
    return %12 : tensor<1x64xbf16> loc(#loc5836)
  } loc(#loc5836)
  func.func private @_einsum_652(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5852)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2195)
  } loc(#loc2195)
  func.func private @_where_653(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5853)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5854)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5855)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5856)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5857)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2204)
  } loc(#loc2204)
  func.func private @_einsum_654(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5858)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5859)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2215)
  } loc(#loc2215)
  func.func private @_einsum_655(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5860)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2216)
  } loc(#loc2216)
  func.func private @_einsum_656(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5861)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2236)
  } loc(#loc2236)
  func.func private @_einsum_657(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5862)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2248)
  } loc(#loc2248)
  func.func private @_einsum_658(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5863)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5864)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2295)
  } loc(#loc2295)
  func.func private @softplus_659(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2301)
    %1 = call @logaddexp_660(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5865)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2301)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5866)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5866)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2301)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5867)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5868)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5869)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2301)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5866)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5866)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2301)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5867)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5868)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5869)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5870)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5871)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2301)
  } loc(#loc2301)
  func.func private @logaddexp_660(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5872)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5872)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5873)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5873)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5874)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5875)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5875)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5876)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5877)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5878)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5879)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5875)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5880)
    return %12 : tensor<1x64xbf16> loc(#loc5865)
  } loc(#loc5865)
  func.func private @_einsum_661(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5881)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2305)
  } loc(#loc2305)
  func.func private @_where_662(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5882)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5883)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5884)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5885)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5886)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2314)
  } loc(#loc2314)
  func.func private @_einsum_663(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5887)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5888)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2325)
  } loc(#loc2325)
  func.func private @_einsum_664(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5889)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2326)
  } loc(#loc2326)
  func.func private @_einsum_665(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5890)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2346)
  } loc(#loc2346)
  func.func private @_einsum_666(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5891)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2358)
  } loc(#loc2358)
  func.func private @_einsum_667(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5892)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5893)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2405)
  } loc(#loc2405)
  func.func private @softplus_668(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2411)
    %1 = call @logaddexp_669(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5894)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2411)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5895)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5895)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2411)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5896)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5897)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5898)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2411)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5895)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5895)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2411)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5896)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5897)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5898)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5899)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5900)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2411)
  } loc(#loc2411)
  func.func private @logaddexp_669(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5901)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5901)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5902)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5902)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5903)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5904)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5904)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5905)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5906)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5907)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5908)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5904)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5909)
    return %12 : tensor<1x64xbf16> loc(#loc5894)
  } loc(#loc5894)
  func.func private @_einsum_670(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5910)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2415)
  } loc(#loc2415)
  func.func private @_where_671(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5911)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5912)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5913)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5914)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5915)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2424)
  } loc(#loc2424)
  func.func private @_einsum_672(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5916)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5917)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2435)
  } loc(#loc2435)
  func.func private @_einsum_673(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5918)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2436)
  } loc(#loc2436)
  func.func private @_einsum_674(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5919)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2456)
  } loc(#loc2456)
  func.func private @_einsum_675(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5920)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2468)
  } loc(#loc2468)
  func.func private @_einsum_676(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5921)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5922)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2515)
  } loc(#loc2515)
  func.func private @softplus_677(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2521)
    %1 = call @logaddexp_678(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5923)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2521)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5924)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5924)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2521)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5925)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5926)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5927)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2521)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5924)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5924)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2521)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5925)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5926)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5927)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5928)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5929)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2521)
  } loc(#loc2521)
  func.func private @logaddexp_678(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5930)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5930)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5931)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5931)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5932)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5933)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5933)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5934)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5935)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5936)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5937)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5933)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5938)
    return %12 : tensor<1x64xbf16> loc(#loc5923)
  } loc(#loc5923)
  func.func private @_einsum_679(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5939)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2525)
  } loc(#loc2525)
  func.func private @_where_680(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5940)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5941)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5942)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5943)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5944)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2534)
  } loc(#loc2534)
  func.func private @_einsum_681(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5945)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5946)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2545)
  } loc(#loc2545)
  func.func private @_einsum_682(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5947)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2546)
  } loc(#loc2546)
  func.func private @_einsum_683(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5948)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2566)
  } loc(#loc2566)
  func.func private @_einsum_684(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5949)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2578)
  } loc(#loc2578)
  func.func private @_einsum_685(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5950)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5951)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2625)
  } loc(#loc2625)
  func.func private @softplus_686(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2631)
    %1 = call @logaddexp_687(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5952)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2631)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5953)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5953)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2631)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5954)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5955)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5956)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2631)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5953)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5953)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2631)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5954)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5955)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5956)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5957)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5958)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2631)
  } loc(#loc2631)
  func.func private @logaddexp_687(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5959)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5959)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5960)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5960)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5961)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5962)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5962)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5963)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5964)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5965)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5966)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5962)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5967)
    return %12 : tensor<1x64xbf16> loc(#loc5952)
  } loc(#loc5952)
  func.func private @_einsum_688(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5968)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2635)
  } loc(#loc2635)
  func.func private @_where_689(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5969)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5970)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc5971)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc5972)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5973)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2644)
  } loc(#loc2644)
  func.func private @_einsum_690(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc5974)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5975)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2655)
  } loc(#loc2655)
  func.func private @_einsum_691(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5976)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2656)
  } loc(#loc2656)
  func.func private @_einsum_692(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc5977)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2676)
  } loc(#loc2676)
  func.func private @_einsum_693(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc5978)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2688)
  } loc(#loc2688)
  func.func private @_einsum_694(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5979)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc5980)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2735)
  } loc(#loc2735)
  func.func private @softplus_695(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2741)
    %1 = call @logaddexp_696(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5981)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2741)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5982)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5982)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2741)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5983)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5984)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5985)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2741)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5982)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5982)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2741)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc5983)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc5984)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5985)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc5986)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc5987)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2741)
  } loc(#loc2741)
  func.func private @logaddexp_696(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5988)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc5988)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5989)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc5989)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc5990)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc5991)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc5991)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc5992)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc5993)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc5994)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc5995)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc5991)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5996)
    return %12 : tensor<1x64xbf16> loc(#loc5981)
  } loc(#loc5981)
  func.func private @_einsum_697(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc5997)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2745)
  } loc(#loc2745)
  func.func private @_where_698(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc5998)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc5999)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6000)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6001)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6002)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2754)
  } loc(#loc2754)
  func.func private @_einsum_699(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6003)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6004)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2765)
  } loc(#loc2765)
  func.func private @_einsum_700(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6005)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2766)
  } loc(#loc2766)
  func.func private @_einsum_701(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6006)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2786)
  } loc(#loc2786)
  func.func private @_einsum_702(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6007)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2798)
  } loc(#loc2798)
  func.func private @_einsum_703(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6008)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6009)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2845)
  } loc(#loc2845)
  func.func private @softplus_704(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2851)
    %1 = call @logaddexp_705(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6010)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2851)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6011)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6011)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2851)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6012)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6013)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6014)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2851)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6011)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6011)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2851)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6012)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6013)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6014)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6015)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6016)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2851)
  } loc(#loc2851)
  func.func private @logaddexp_705(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6017)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6017)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6018)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6018)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6019)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6020)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6020)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6021)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6022)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6023)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6024)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6020)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6025)
    return %12 : tensor<1x64xbf16> loc(#loc6010)
  } loc(#loc6010)
  func.func private @_einsum_706(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6026)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2855)
  } loc(#loc2855)
  func.func private @_where_707(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6027)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6028)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6029)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6030)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6031)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2864)
  } loc(#loc2864)
  func.func private @_einsum_708(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6032)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6033)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2875)
  } loc(#loc2875)
  func.func private @_einsum_709(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6034)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2876)
  } loc(#loc2876)
  func.func private @_einsum_710(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6035)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc2896)
  } loc(#loc2896)
  func.func private @_einsum_711(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6036)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2908)
  } loc(#loc2908)
  func.func private @_einsum_712(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6037)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6038)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc2955)
  } loc(#loc2955)
  func.func private @softplus_713(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2961)
    %1 = call @logaddexp_714(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6039)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2961)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6040)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6040)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2961)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6041)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6042)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6043)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc2961)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6040)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6040)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc2961)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6041)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6042)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6043)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6044)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6045)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc2961)
  } loc(#loc2961)
  func.func private @logaddexp_714(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6046)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6046)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6047)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6047)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6048)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6049)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6049)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6050)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6051)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6052)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6053)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6049)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6054)
    return %12 : tensor<1x64xbf16> loc(#loc6039)
  } loc(#loc6039)
  func.func private @_einsum_715(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6055)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc2965)
  } loc(#loc2965)
  func.func private @_where_716(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6057)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6058)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6059)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6060)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc2974)
  } loc(#loc2974)
  func.func private @_einsum_717(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6061)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6062)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc2985)
  } loc(#loc2985)
  func.func private @_einsum_718(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6063)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc2986)
  } loc(#loc2986)
  func.func private @_einsum_719(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6064)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3006)
  } loc(#loc3006)
  func.func private @_einsum_720(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6065)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3018)
  } loc(#loc3018)
  func.func private @_einsum_721(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6066)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6067)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3065)
  } loc(#loc3065)
  func.func private @softplus_722(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3071)
    %1 = call @logaddexp_723(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6068)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3071)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6069)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6069)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3071)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6070)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6071)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6072)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3071)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6069)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6069)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3071)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6070)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6071)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6072)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6073)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6074)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3071)
  } loc(#loc3071)
  func.func private @logaddexp_723(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6075)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6075)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6076)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6076)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6077)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6078)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6078)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6079)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6080)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6081)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6082)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6078)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6083)
    return %12 : tensor<1x64xbf16> loc(#loc6068)
  } loc(#loc6068)
  func.func private @_einsum_724(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6084)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3075)
  } loc(#loc3075)
  func.func private @_where_725(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6085)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6086)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6087)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6088)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6089)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3084)
  } loc(#loc3084)
  func.func private @_einsum_726(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6090)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6091)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3095)
  } loc(#loc3095)
  func.func private @_einsum_727(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6092)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3096)
  } loc(#loc3096)
  func.func private @_einsum_728(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6093)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3116)
  } loc(#loc3116)
  func.func private @_einsum_729(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6094)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3128)
  } loc(#loc3128)
  func.func private @_einsum_730(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6095)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6096)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3175)
  } loc(#loc3175)
  func.func private @softplus_731(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3181)
    %1 = call @logaddexp_732(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6097)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3181)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6098)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6098)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3181)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6099)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6100)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6101)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3181)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6098)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6098)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3181)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6099)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6100)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6101)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6102)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6103)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3181)
  } loc(#loc3181)
  func.func private @logaddexp_732(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6104)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6104)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6105)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6105)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6106)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6107)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6107)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6108)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6109)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6110)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6111)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6107)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6112)
    return %12 : tensor<1x64xbf16> loc(#loc6097)
  } loc(#loc6097)
  func.func private @_einsum_733(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6113)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3185)
  } loc(#loc3185)
  func.func private @_where_734(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6114)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6115)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6116)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6117)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6118)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3194)
  } loc(#loc3194)
  func.func private @_einsum_735(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6119)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6120)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3205)
  } loc(#loc3205)
  func.func private @_einsum_736(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6121)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3206)
  } loc(#loc3206)
  func.func private @_einsum_737(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6122)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3226)
  } loc(#loc3226)
  func.func private @_einsum_738(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6123)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3238)
  } loc(#loc3238)
  func.func private @_einsum_739(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6124)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6125)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3285)
  } loc(#loc3285)
  func.func private @softplus_740(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3291)
    %1 = call @logaddexp_741(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6126)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3291)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6127)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6127)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3291)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6128)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6129)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6130)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3291)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6127)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6127)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3291)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6128)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6129)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6130)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6131)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6132)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3291)
  } loc(#loc3291)
  func.func private @logaddexp_741(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6133)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6133)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6134)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6134)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6135)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6136)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6136)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6137)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6138)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6139)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6140)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6136)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6141)
    return %12 : tensor<1x64xbf16> loc(#loc6126)
  } loc(#loc6126)
  func.func private @_einsum_742(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6142)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3295)
  } loc(#loc3295)
  func.func private @_where_743(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6143)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6144)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6145)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6146)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6147)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3304)
  } loc(#loc3304)
  func.func private @_einsum_744(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6148)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6149)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3315)
  } loc(#loc3315)
  func.func private @_einsum_745(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6150)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3316)
  } loc(#loc3316)
  func.func private @_einsum_746(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6151)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3336)
  } loc(#loc3336)
  func.func private @_einsum_747(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6152)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3348)
  } loc(#loc3348)
  func.func private @_einsum_748(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6153)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6154)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3395)
  } loc(#loc3395)
  func.func private @softplus_749(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3401)
    %1 = call @logaddexp_750(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6155)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3401)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6156)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6156)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3401)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6157)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6158)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6159)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3401)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6156)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6156)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3401)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6157)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6158)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6159)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6160)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6161)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3401)
  } loc(#loc3401)
  func.func private @logaddexp_750(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6162)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6162)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6163)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6163)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6164)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6165)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6165)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6166)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6167)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6168)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6169)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6165)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6170)
    return %12 : tensor<1x64xbf16> loc(#loc6155)
  } loc(#loc6155)
  func.func private @_einsum_751(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6171)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3405)
  } loc(#loc3405)
  func.func private @_where_752(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6172)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6173)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6174)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6175)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6176)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3414)
  } loc(#loc3414)
  func.func private @_einsum_753(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6177)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6178)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3425)
  } loc(#loc3425)
  func.func private @_einsum_754(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6179)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3426)
  } loc(#loc3426)
  func.func private @_einsum_755(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6180)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3446)
  } loc(#loc3446)
  func.func private @_einsum_756(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6181)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3458)
  } loc(#loc3458)
  func.func private @_einsum_757(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6182)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6183)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3505)
  } loc(#loc3505)
  func.func private @softplus_758(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3511)
    %1 = call @logaddexp_759(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6184)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3511)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6185)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6185)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3511)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6186)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6187)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6188)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3511)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6185)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6185)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3511)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6186)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6187)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6188)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6189)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6190)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3511)
  } loc(#loc3511)
  func.func private @logaddexp_759(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6191)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6191)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6192)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6192)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6193)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6194)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6194)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6195)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6196)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6197)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6198)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6194)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6199)
    return %12 : tensor<1x64xbf16> loc(#loc6184)
  } loc(#loc6184)
  func.func private @_einsum_760(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6200)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3515)
  } loc(#loc3515)
  func.func private @_where_761(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6201)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6202)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6203)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6204)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6205)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3524)
  } loc(#loc3524)
  func.func private @_einsum_762(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6206)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6207)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3535)
  } loc(#loc3535)
  func.func private @_einsum_763(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6208)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3536)
  } loc(#loc3536)
  func.func private @_einsum_764(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6209)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3556)
  } loc(#loc3556)
  func.func private @_einsum_765(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6210)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3568)
  } loc(#loc3568)
  func.func private @_einsum_766(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6211)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6212)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3615)
  } loc(#loc3615)
  func.func private @softplus_767(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3621)
    %1 = call @logaddexp_768(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6213)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3621)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6214)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6214)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3621)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6215)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6216)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6217)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3621)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6214)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6214)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3621)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6215)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6216)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6217)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6218)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6219)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3621)
  } loc(#loc3621)
  func.func private @logaddexp_768(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6220)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6220)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6221)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6221)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6222)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6223)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6223)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6224)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6225)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6226)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6227)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6223)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6228)
    return %12 : tensor<1x64xbf16> loc(#loc6213)
  } loc(#loc6213)
  func.func private @_einsum_769(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6229)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3625)
  } loc(#loc3625)
  func.func private @_where_770(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6230)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6231)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6232)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6233)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6234)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3634)
  } loc(#loc3634)
  func.func private @_einsum_771(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6235)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6236)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3645)
  } loc(#loc3645)
  func.func private @_einsum_772(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6237)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3646)
  } loc(#loc3646)
  func.func private @_einsum_773(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6238)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3666)
  } loc(#loc3666)
  func.func private @_einsum_774(%arg0: tensor<1x1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6239)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3678)
  } loc(#loc3678)
  func.func private @_einsum_775(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0] x [0], contracting_dims = [2] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x3x2048x32x64xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6240)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x32x64x1x2048xbf16>) -> tensor<1x3x1x2048x32x64xbf16> loc(#loc6241)
    return %1 : tensor<1x3x1x2048x32x64xbf16> loc(#loc3725)
  } loc(#loc3725)
  func.func private @softplus_776(%arg0: tensor<1x64xbf16> loc(unknown)) -> (tensor<1x64xbf16>, tensor<1x64xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3731)
    %1 = call @logaddexp_777(%arg0, %0) : (tensor<1x64xbf16>, tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6242)
    %2 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3731)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6243)
    %4 = stablehlo.compare  EQ, %arg0, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6243)
    %5 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3731)
    %6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6244)
    %7 = stablehlo.broadcast_in_dim %6, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6245)
    %8 = stablehlo.select %4, %7, %arg0 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6246)
    %9 = stablehlo.constant dense<0x7F80> : tensor<bf16> loc(#loc3731)
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6243)
    %11 = stablehlo.compare  EQ, %1, %10,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6243)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc3731)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<bf16>) -> tensor<64xbf16> loc(#loc6244)
    %14 = stablehlo.broadcast_in_dim %13, dims = [1] : (tensor<64xbf16>) -> tensor<1x64xbf16> loc(#loc6245)
    %15 = stablehlo.select %11, %14, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6246)
    %16 = stablehlo.subtract %8, %15 : tensor<1x64xbf16> loc(#loc6247)
    %17 = stablehlo.exponential %16 : tensor<1x64xbf16> loc(#loc6248)
    return %1, %17 : tensor<1x64xbf16>, tensor<1x64xbf16> loc(#loc3731)
  } loc(#loc3731)
  func.func private @logaddexp_777(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<bf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6249)
    %1 = stablehlo.maximum %arg0, %0 : tensor<1x64xbf16> loc(#loc6249)
    %2 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6250)
    %3 = stablehlo.subtract %arg0, %2 : tensor<1x64xbf16> loc(#loc6250)
    %4 = stablehlo.compare  NE, %3, %3,  FLOAT : (tensor<1x64xbf16>, tensor<1x64xbf16>) -> tensor<1x64xi1> loc(#loc6251)
    %5 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6252)
    %6 = stablehlo.add %arg0, %5 : tensor<1x64xbf16> loc(#loc6252)
    %7 = stablehlo.abs %3 : tensor<1x64xbf16> loc(#loc6253)
    %8 = stablehlo.negate %7 : tensor<1x64xbf16> loc(#loc6254)
    %9 = stablehlo.exponential %8 : tensor<1x64xbf16> loc(#loc6255)
    %10 = stablehlo.log_plus_one %9 : tensor<1x64xbf16> loc(#loc6256)
    %11 = stablehlo.add %1, %10 : tensor<1x64xbf16> loc(#loc6252)
    %12 = stablehlo.select %4, %6, %11 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc6257)
    return %12 : tensor<1x64xbf16> loc(#loc6242)
  } loc(#loc6242)
  func.func private @_einsum_778(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x32x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1, 3] x [0, 1, 3], contracting_dims = [4] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6258)
    return %0 : tensor<1x1x32x2048x2048xbf16> loc(#loc3735)
  } loc(#loc3735)
  func.func private @_where_779(%arg0: tensor<1x1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> (tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1>) {
    %0 = stablehlo.reshape %arg0 : (tensor<1x1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6259)
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 3, 4] : (tensor<1x1x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6260)
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1x32x2048x2048xf32> loc(#loc6261)
    %3 = stablehlo.broadcast_in_dim %2, dims = [1, 2, 3, 4] : (tensor<1x32x2048x2048xf32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6262)
    %4 = stablehlo.select %1, %arg1, %3 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6263)
    return %4, %1 : tensor<1x1x32x2048x2048xf32>, tensor<1x1x32x2048x2048xi1> loc(#loc3744)
  } loc(#loc3744)
  func.func private @_einsum_780(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x32x64xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, batching_dims = [0, 1, 3] x [0, 1, 2], contracting_dims = [2] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6264)
    %1 = stablehlo.transpose %0, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6265)
    return %1 : tensor<1x1x2048x32x64xbf16> loc(#loc3755)
  } loc(#loc3755)
  func.func private @_einsum_781(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [4, 3] x [3, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6266)
    return %0 : tensor<1x1x2048x2048xbf16> loc(#loc3756)
  } loc(#loc3756)
  func.func private @_einsum_782(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x1x2048x8192xbf16> {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6267)
    return %0 : tensor<1x1x2048x8192xbf16> loc(#loc3776)
  } loc(#loc3776)
  func.func private @_where_783(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_784(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_785(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_786(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_787(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_788(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_789(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_790(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_791(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_792(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_793(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_794(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_795(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_796(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_797(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_798(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_799(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_800(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_801(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_802(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_803(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_804(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_805(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_806(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_807(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_808(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_809(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_810(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_811(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_812(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_813(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_814(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_815(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_816(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_817(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_818(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_819(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_820(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_821(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_822(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_823(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_824(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_825(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_826(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_827(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_828(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_829(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_where_830(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc3789)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6268)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc3789)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6269)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1> loc(#loc6269)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1xi1>, tensor<1xf32> loc(#loc6270)
    return %5 : tensor<1xf32> loc(#loc3789)
  } loc(#loc3789)
  func.func private @_pad_831(%arg0: tensor<2x1x2048x2048xbf16> loc(unknown)) -> tensor<1x1x2048x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %1 = stablehlo.pad %arg0, %0, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %2 = "stablehlo.slice"(%1) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %3 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %4 = stablehlo.pad %2, %3, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %5 = "stablehlo.slice"(%4) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %6 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %7 = stablehlo.pad %5, %6, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %8 = "stablehlo.slice"(%7) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %9 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %10 = stablehlo.pad %8, %9, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %11 = "stablehlo.slice"(%10) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %13 = stablehlo.pad %11, %12, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %14 = "stablehlo.slice"(%13) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %15 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %16 = stablehlo.pad %14, %15, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %17 = "stablehlo.slice"(%16) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %18 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %19 = stablehlo.pad %17, %18, low = [0, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6271)
    %20 = "stablehlo.slice"(%19) {limit_indices = dense<[2, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<2x1x2048x2048xbf16>) -> tensor<2x1x2048x2048xbf16> loc(#loc6272)
    %21 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4048)
    %22 = stablehlo.pad %20, %21, low = [-1, 0, 0, 0], high = [0, 0, 0, 0], interior = [0, 0, 0, 0] : (tensor<2x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6273)
    %23 = "stablehlo.slice"(%22) {limit_indices = dense<[1, 1, 2048, 2048]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>} : (tensor<1x1x2048x2048xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6274)
    return %23 : tensor<1x1x2048x2048xbf16> loc(#loc4048)
  } loc(#loc4048)
  func.func private @_where_832(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_833(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6277)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6278)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6279)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4058)
  } loc(#loc4058)
  func.func private @_where_834(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_835(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_836(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6282)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6283)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6284)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4064)
  } loc(#loc4064)
  func.func private @_where_837(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_838(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_839(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_840(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6286)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6287)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4084)
  } loc(#loc4084)
  func.func private @_where_841(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_842(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6289)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6290)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6291)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6265)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4085)
  } loc(#loc4085)
  func.func private @_where_843(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4095)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6292)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4095)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6293)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6293)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6263)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4095)
  } loc(#loc4095)
  func.func private @_einsum_844(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6294)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6295)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6296)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6295)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4098)
  } loc(#loc4098)
  func.func private @softplus_845(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6297)
    return %0 : tensor<1x64xbf16> loc(#loc4105)
  } loc(#loc4105)
  func.func private @_where_846(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_847(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6241)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6299)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6300)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6301)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4107)
  } loc(#loc4107)
  func.func private @_where_848(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_849(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_850(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_851(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_852(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6303)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6304)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6305)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4130)
  } loc(#loc4130)
  func.func private @_where_853(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_854(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_855(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6306)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6307)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6308)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4136)
  } loc(#loc4136)
  func.func private @_where_856(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_857(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_858(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_859(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6309)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6310)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4156)
  } loc(#loc4156)
  func.func private @_where_860(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_861(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6311)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6312)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6313)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6236)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4157)
  } loc(#loc4157)
  func.func private @_where_862(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4167)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6314)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4167)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6315)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6315)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6234)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4167)
  } loc(#loc4167)
  func.func private @_einsum_863(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6316)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6317)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6318)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6317)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4170)
  } loc(#loc4170)
  func.func private @softplus_864(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6319)
    return %0 : tensor<1x64xbf16> loc(#loc4177)
  } loc(#loc4177)
  func.func private @_where_865(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_866(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6212)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6320)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6321)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6322)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4179)
  } loc(#loc4179)
  func.func private @_where_867(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_868(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_869(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_870(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_871(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6323)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6324)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6325)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4202)
  } loc(#loc4202)
  func.func private @_where_872(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_873(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_874(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6326)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6327)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6328)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4208)
  } loc(#loc4208)
  func.func private @_where_875(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_876(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_877(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_878(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6329)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6330)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4228)
  } loc(#loc4228)
  func.func private @_where_879(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_880(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6331)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6332)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6333)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6207)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4229)
  } loc(#loc4229)
  func.func private @_where_881(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4239)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6334)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4239)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6335)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6335)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6205)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4239)
  } loc(#loc4239)
  func.func private @_einsum_882(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6336)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6337)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6338)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6337)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4242)
  } loc(#loc4242)
  func.func private @softplus_883(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6339)
    return %0 : tensor<1x64xbf16> loc(#loc4249)
  } loc(#loc4249)
  func.func private @_where_884(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_885(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6183)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6340)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6341)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6342)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4251)
  } loc(#loc4251)
  func.func private @_where_886(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_887(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_888(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_889(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_890(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6343)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6344)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6345)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4274)
  } loc(#loc4274)
  func.func private @_where_891(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_892(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_893(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6346)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6347)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6348)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4280)
  } loc(#loc4280)
  func.func private @_where_894(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_895(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_896(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_897(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6349)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6350)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4300)
  } loc(#loc4300)
  func.func private @_where_898(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_899(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6351)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6352)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6353)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6178)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4301)
  } loc(#loc4301)
  func.func private @_where_900(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4311)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6354)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4311)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6355)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6355)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6176)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4311)
  } loc(#loc4311)
  func.func private @_einsum_901(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6356)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6357)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6358)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6357)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4314)
  } loc(#loc4314)
  func.func private @softplus_902(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6359)
    return %0 : tensor<1x64xbf16> loc(#loc4321)
  } loc(#loc4321)
  func.func private @_where_903(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_904(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6154)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6360)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6361)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6362)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4323)
  } loc(#loc4323)
  func.func private @_where_905(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_906(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_907(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_908(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_909(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6363)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6364)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6365)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4346)
  } loc(#loc4346)
  func.func private @_where_910(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_911(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_912(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6366)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6367)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6368)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4352)
  } loc(#loc4352)
  func.func private @_where_913(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_914(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_915(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_916(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6369)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6370)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4372)
  } loc(#loc4372)
  func.func private @_where_917(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_918(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6371)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6372)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6373)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6149)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4373)
  } loc(#loc4373)
  func.func private @_where_919(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4383)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6374)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4383)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6375)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6375)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6147)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4383)
  } loc(#loc4383)
  func.func private @_einsum_920(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6376)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6377)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6378)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6377)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4386)
  } loc(#loc4386)
  func.func private @softplus_921(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6379)
    return %0 : tensor<1x64xbf16> loc(#loc4393)
  } loc(#loc4393)
  func.func private @_where_922(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_923(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6125)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6380)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6381)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6382)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4395)
  } loc(#loc4395)
  func.func private @_where_924(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_925(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_926(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_927(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_928(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6383)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6384)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6385)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4418)
  } loc(#loc4418)
  func.func private @_where_929(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_930(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_931(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6386)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6387)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6388)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4424)
  } loc(#loc4424)
  func.func private @_where_932(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_933(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_934(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_935(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6389)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6390)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4444)
  } loc(#loc4444)
  func.func private @_where_936(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_937(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6391)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6392)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6393)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6120)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4445)
  } loc(#loc4445)
  func.func private @_where_938(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4455)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6394)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4455)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6395)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6395)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6118)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4455)
  } loc(#loc4455)
  func.func private @_einsum_939(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6396)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6397)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6398)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6397)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4458)
  } loc(#loc4458)
  func.func private @softplus_940(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6399)
    return %0 : tensor<1x64xbf16> loc(#loc4465)
  } loc(#loc4465)
  func.func private @_where_941(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_942(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6096)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6400)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6401)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6402)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4467)
  } loc(#loc4467)
  func.func private @_where_943(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_944(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_945(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_946(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_947(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6403)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6404)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6405)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4490)
  } loc(#loc4490)
  func.func private @_where_948(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_949(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_950(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6406)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6407)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6408)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4496)
  } loc(#loc4496)
  func.func private @_where_951(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_952(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_953(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_954(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6409)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6410)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4516)
  } loc(#loc4516)
  func.func private @_where_955(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_956(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6411)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6412)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6413)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6091)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4517)
  } loc(#loc4517)
  func.func private @_where_957(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4527)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6414)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4527)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6415)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6415)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6089)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4527)
  } loc(#loc4527)
  func.func private @_einsum_958(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6416)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6417)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6418)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6417)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4530)
  } loc(#loc4530)
  func.func private @softplus_959(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6419)
    return %0 : tensor<1x64xbf16> loc(#loc4537)
  } loc(#loc4537)
  func.func private @_where_960(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_961(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6067)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6420)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6421)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6422)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4539)
  } loc(#loc4539)
  func.func private @_where_962(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_963(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_964(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_965(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_966(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6423)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6424)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6425)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4562)
  } loc(#loc4562)
  func.func private @_where_967(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_968(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_969(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6426)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6427)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6428)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4568)
  } loc(#loc4568)
  func.func private @_where_970(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_971(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_972(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_973(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6429)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6430)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4588)
  } loc(#loc4588)
  func.func private @_where_974(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_975(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6431)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6432)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6433)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6062)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4589)
  } loc(#loc4589)
  func.func private @_where_976(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4599)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6434)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4599)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6435)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6435)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6060)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4599)
  } loc(#loc4599)
  func.func private @_einsum_977(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6436)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6437)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6438)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6437)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4602)
  } loc(#loc4602)
  func.func private @softplus_978(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6439)
    return %0 : tensor<1x64xbf16> loc(#loc4609)
  } loc(#loc4609)
  func.func private @_where_979(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_980(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6038)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6440)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6441)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6442)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4611)
  } loc(#loc4611)
  func.func private @_where_981(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_982(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_983(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_984(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_985(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6443)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6444)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6445)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4634)
  } loc(#loc4634)
  func.func private @_where_986(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_987(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_988(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6446)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6447)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6448)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4640)
  } loc(#loc4640)
  func.func private @_where_989(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_990(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_991(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_992(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6449)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6450)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4660)
  } loc(#loc4660)
  func.func private @_where_993(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_994(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6451)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6452)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6453)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6033)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4661)
  } loc(#loc4661)
  func.func private @_where_995(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4671)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6454)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4671)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6455)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6455)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6031)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4671)
  } loc(#loc4671)
  func.func private @_einsum_996(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6456)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6457)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6458)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6457)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4674)
  } loc(#loc4674)
  func.func private @softplus_997(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6459)
    return %0 : tensor<1x64xbf16> loc(#loc4681)
  } loc(#loc4681)
  func.func private @_where_998(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_999(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc6009)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6460)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6461)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6462)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4683)
  } loc(#loc4683)
  func.func private @_where_1000(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1001(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1002(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1003(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1004(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6463)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6464)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6465)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4706)
  } loc(#loc4706)
  func.func private @_where_1005(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1006(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1007(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6466)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6467)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6468)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4712)
  } loc(#loc4712)
  func.func private @_where_1008(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1009(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1010(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1011(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6469)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6470)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4732)
  } loc(#loc4732)
  func.func private @_where_1012(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1013(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6471)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6472)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6473)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6004)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4733)
  } loc(#loc4733)
  func.func private @_where_1014(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4743)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6474)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4743)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6475)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6475)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc6002)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4743)
  } loc(#loc4743)
  func.func private @_einsum_1015(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6476)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6477)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6478)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6477)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4746)
  } loc(#loc4746)
  func.func private @softplus_1016(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6479)
    return %0 : tensor<1x64xbf16> loc(#loc4753)
  } loc(#loc4753)
  func.func private @_where_1017(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1018(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5980)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6480)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6481)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6482)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4755)
  } loc(#loc4755)
  func.func private @_where_1019(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1020(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1021(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1022(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1023(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6483)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6484)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6485)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4778)
  } loc(#loc4778)
  func.func private @_where_1024(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1025(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1026(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6486)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6487)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6488)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4784)
  } loc(#loc4784)
  func.func private @_where_1027(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1028(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1029(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1030(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6489)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6490)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4804)
  } loc(#loc4804)
  func.func private @_where_1031(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1032(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6491)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6492)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6493)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5975)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4805)
  } loc(#loc4805)
  func.func private @_where_1033(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4815)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6494)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4815)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6495)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6495)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5973)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4815)
  } loc(#loc4815)
  func.func private @_einsum_1034(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6496)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6497)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6498)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6497)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4818)
  } loc(#loc4818)
  func.func private @softplus_1035(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6499)
    return %0 : tensor<1x64xbf16> loc(#loc4825)
  } loc(#loc4825)
  func.func private @_where_1036(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1037(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5951)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6500)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6501)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6502)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4827)
  } loc(#loc4827)
  func.func private @_where_1038(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1039(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1040(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1041(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1042(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6503)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6504)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6505)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4850)
  } loc(#loc4850)
  func.func private @_where_1043(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1044(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1045(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6506)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6507)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6508)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4856)
  } loc(#loc4856)
  func.func private @_where_1046(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1047(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1048(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1049(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6509)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6510)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4876)
  } loc(#loc4876)
  func.func private @_where_1050(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1051(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6511)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6512)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6513)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5946)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4877)
  } loc(#loc4877)
  func.func private @_where_1052(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4887)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6514)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4887)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6515)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6515)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5944)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4887)
  } loc(#loc4887)
  func.func private @_einsum_1053(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6516)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6517)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6518)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6517)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4890)
  } loc(#loc4890)
  func.func private @softplus_1054(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6519)
    return %0 : tensor<1x64xbf16> loc(#loc4897)
  } loc(#loc4897)
  func.func private @_where_1055(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1056(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5922)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6520)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6521)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6522)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4899)
  } loc(#loc4899)
  func.func private @_where_1057(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1058(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1059(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1060(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1061(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6523)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6524)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6525)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4922)
  } loc(#loc4922)
  func.func private @_where_1062(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1063(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1064(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6526)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6527)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6528)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc4928)
  } loc(#loc4928)
  func.func private @_where_1065(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1066(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1067(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1068(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6529)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6530)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc4948)
  } loc(#loc4948)
  func.func private @_where_1069(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1070(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6531)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6532)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6533)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5917)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4949)
  } loc(#loc4949)
  func.func private @_where_1071(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc4959)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6534)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4959)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6535)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6535)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5915)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc4959)
  } loc(#loc4959)
  func.func private @_einsum_1072(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6536)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6537)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6538)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6537)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc4962)
  } loc(#loc4962)
  func.func private @softplus_1073(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6539)
    return %0 : tensor<1x64xbf16> loc(#loc4969)
  } loc(#loc4969)
  func.func private @_where_1074(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1075(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5893)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6540)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6541)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6542)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc4971)
  } loc(#loc4971)
  func.func private @_where_1076(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1077(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1078(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1079(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1080(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6543)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6544)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6545)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc4994)
  } loc(#loc4994)
  func.func private @_where_1081(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1082(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1083(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6546)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6547)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6548)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc5000)
  } loc(#loc5000)
  func.func private @_where_1084(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1085(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1086(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1087(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6549)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6550)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc5020)
  } loc(#loc5020)
  func.func private @_where_1088(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1089(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6551)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6552)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6553)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5888)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc5021)
  } loc(#loc5021)
  func.func private @_where_1090(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5031)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6554)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc5031)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6555)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6555)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5886)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc5031)
  } loc(#loc5031)
  func.func private @_einsum_1091(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6556)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6557)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6558)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6557)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc5034)
  } loc(#loc5034)
  func.func private @softplus_1092(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6559)
    return %0 : tensor<1x64xbf16> loc(#loc5041)
  } loc(#loc5041)
  func.func private @_where_1093(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1094(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5864)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6560)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6561)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6562)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc5043)
  } loc(#loc5043)
  func.func private @_where_1095(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1096(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1097(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1098(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1099(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6563)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6564)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6565)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc5066)
  } loc(#loc5066)
  func.func private @_where_1100(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1101(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1102(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6566)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6567)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6568)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc5072)
  } loc(#loc5072)
  func.func private @_where_1103(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1104(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1105(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1106(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6569)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6570)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc5092)
  } loc(#loc5092)
  func.func private @_where_1107(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1108(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6571)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6572)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6573)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5859)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc5093)
  } loc(#loc5093)
  func.func private @_where_1109(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5103)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6574)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc5103)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6575)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6575)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5857)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc5103)
  } loc(#loc5103)
  func.func private @_einsum_1110(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6576)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6577)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6578)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6577)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc5106)
  } loc(#loc5106)
  func.func private @softplus_1111(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6579)
    return %0 : tensor<1x64xbf16> loc(#loc5113)
  } loc(#loc5113)
  func.func private @_where_1112(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1113(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5835)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6580)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6581)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6582)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc5115)
  } loc(#loc5115)
  func.func private @_where_1114(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1115(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1116(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1117(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1118(%arg0: tensor<1x8192x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x8192xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x8192xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6583)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x2048x8192xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6584)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x8192x2048xbf16>) -> tensor<1x1x2048x8192xbf16> loc(#loc6585)
    return %2, %1 : tensor<1x1x2048x8192xbf16>, tensor<1x8192x2048xbf16> loc(#loc5138)
  } loc(#loc5138)
  func.func private @_where_1119(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<1x8192x2048xbf16> loc(unknown)) -> tensor<1x8192x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192x2048xbf16> loc(#loc6280)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192x2048xi1>, tensor<1x8192x2048xi1>) -> tensor<1x8192x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x8192x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1120(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<1x8192xbf16> loc(unknown)) -> tensor<1x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x8192xbf16> loc(#loc6281)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x8192xi1>, tensor<1x8192xi1>) -> tensor<1x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x8192xi1>, tensor<1x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1121(%arg0: tensor<1x2048x8192xbf16> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown), %arg2: tensor<1x1x2048x8192xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x8192x2048xbf16> loc(#loc6586)
    %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<1x8192x2048xbf16>) -> tensor<1x2048x8192xbf16> loc(#loc6587)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x8192xbf16>, tensor<1x2048x8192xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6588)
    return %2, %1 : tensor<1x1x2048x2048xbf16>, tensor<1x2048x8192xbf16> loc(#loc5144)
  } loc(#loc5144)
  func.func private @_where_1122(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<1x2048x8192xbf16> loc(unknown)) -> tensor<1x2048x8192xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x8192xbf16> loc(#loc6285)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x8192xi1>, tensor<1x2048x8192xi1>) -> tensor<1x2048x8192xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x8192xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1123(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1124(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1125(%arg0: tensor<1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6589)
    %1 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0] x [0], contracting_dims = [3] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x1x2048x2048xbf16>, tensor<1x2048x32x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6590)
    return %1, %0 : tensor<1x1x2048x32x64xbf16>, tensor<1x2048x32x64xbf16> loc(#loc5164)
  } loc(#loc5164)
  func.func private @_where_1126(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<1x2048x32x64xbf16> loc(unknown)) -> tensor<1x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048x32x64xbf16> loc(#loc6288)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xi1>) -> tensor<1x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1127(%arg0: tensor<1x1x32x2048x2048xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 3, 4, 2] : (tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6591)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [4], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x2048xbf16> loc(#loc6592)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0, 1, 2] x [0, 1, 2], contracting_dims = [4] x [3], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x64x2048xbf16>, tensor<1x1x32x2048x2048xbf16>) -> tensor<1x1x32x64x2048xbf16> loc(#loc6593)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x1x32x64x2048xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc5830)
    return %1, %3 : tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc5165)
  } loc(#loc5165)
  func.func private @_where_1128(%arg0: tensor<1x1x32x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x32x2048x2048xf32> loc(unknown)) -> tensor<1x1x32x2048x2048xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32> loc(#loc5175)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1x1x32x2048x2048xf32> loc(#loc6594)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc5175)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6595)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xi1>) -> tensor<1x1x32x2048x2048xi1> loc(#loc6595)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x32x2048x2048xi1>, tensor<1x1x32x2048x2048xf32> loc(#loc5828)
    return %5 : tensor<1x1x32x2048x2048xf32> loc(#loc5175)
  } loc(#loc5175)
  func.func private @_einsum_1129(%arg0: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg1: tensor<1x1x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x1x32x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16>) {
    %0 = stablehlo.dot_general %arg2, %arg1, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [3] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6596)
    %1 = stablehlo.transpose %0, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6597)
    %2 = stablehlo.dot_general %arg2, %arg0, batching_dims = [0, 1, 2] x [0, 1, 3], contracting_dims = [4] x [2], precision = [DEFAULT, DEFAULT] : (tensor<1x1x32x2048x2048xbf16>, tensor<1x1x2048x32x64xbf16>) -> tensor<1x1x32x2048x64xbf16> loc(#loc6598)
    %3 = stablehlo.transpose %2, dims = [0, 1, 3, 2, 4] : (tensor<1x1x32x2048x64xbf16>) -> tensor<1x1x2048x32x64xbf16> loc(#loc6597)
    return %3, %1 : tensor<1x1x2048x32x64xbf16>, tensor<1x1x2048x32x64xbf16> loc(#loc5178)
  } loc(#loc5178)
  func.func private @softplus_1130(%arg0: tensor<1x64xbf16> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.multiply %arg1, %arg0 : tensor<1x64xbf16> loc(#loc6599)
    return %0 : tensor<1x64xbf16> loc(#loc5185)
  } loc(#loc5185)
  func.func private @_where_1131(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<1x64xbf16> loc(unknown)) -> tensor<1x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x64xbf16> loc(#loc6298)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x64xi1>, tensor<1x64xi1>) -> tensor<1x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x64xi1>, tensor<1x64xbf16> loc(#loc5798)
    return %5 : tensor<1x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_einsum_1132(%arg0: tensor<1x1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown), %arg2: tensor<1x3x1x2048x32x64xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16>) {
    %0 = stablehlo.transpose %arg2, dims = [0, 1, 4, 5, 2, 3] : (tensor<1x3x1x2048x32x64xbf16>) -> tensor<1x3x32x64x1x2048xbf16> loc(#loc5806)
    %1 = stablehlo.dot_general %0, %arg1, batching_dims = [0] x [0], contracting_dims = [1, 2, 3] x [1, 3, 4], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x3x2048x32x64xbf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6600)
    %2 = stablehlo.dot_general %0, %arg0, batching_dims = [0] x [0], contracting_dims = [4, 5] x [1, 2], precision = [DEFAULT, DEFAULT] : (tensor<1x3x32x64x1x2048xbf16>, tensor<1x1x2048x2048xbf16>) -> tensor<1x3x32x64x2048xbf16> loc(#loc6601)
    %3 = stablehlo.transpose %2, dims = [0, 1, 4, 2, 3] : (tensor<1x3x32x64x2048xbf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6602)
    return %1, %3 : tensor<1x1x2048x2048xbf16>, tensor<1x3x2048x32x64xbf16> loc(#loc5187)
  } loc(#loc5187)
  func.func private @_where_1133(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<1x3x2048x32x64xbf16> loc(unknown)) -> tensor<1x3x2048x32x64xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x3x2048x32x64xbf16> loc(#loc6302)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xi1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xbf16> loc(#loc5798)
    return %5 : tensor<1x3x2048x32x64xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1134(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1135(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<1x2048xbf16> loc(unknown)) -> tensor<1x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc4056)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x2048xbf16> loc(#loc6275)
    %2 = stablehlo.constant dense<true> : tensor<i1> loc(#loc4056)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6276)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x2048xi1>, tensor<1x2048xi1>) -> tensor<1x2048xi1> loc(#loc6276)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x2048xi1>, tensor<1x2048xbf16> loc(#loc5798)
    return %5 : tensor<1x2048xbf16> loc(#loc4056)
  } loc(#loc4056)
  func.func private @_where_1136(%arg0: tensor<1x1x2048x2048xi1> loc(unknown), %arg1: tensor<1x1x2048x2048xbf16> loc(unknown)) -> (tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x2048xbf16>) {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5206)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<bf16>) -> tensor<1x1x2048x2048xbf16> loc(#loc6603)
    %2 = stablehlo.constant dense<false> : tensor<i1> loc(#loc5206)
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc6604)
    %4 = stablehlo.compare  EQ, %arg0, %3,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6604)
    %5 = stablehlo.select %4, %arg1, %1 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xbf16> loc(#loc5796)
    %6 = stablehlo.constant dense<true> : tensor<i1> loc(#loc5206)
    %7 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor<i1>) -> tensor<1x1x2048x2048xi1> loc(#loc6604)
    %8 = stablehlo.compare  EQ, %arg0, %7,  UNSIGNED : (tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xi1>) -> tensor<1x1x2048x2048xi1> loc(#loc6604)
    %9 = stablehlo.select %8, %arg1, %1 : tensor<1x1x2048x2048xi1>, tensor<1x1x2048x2048xbf16> loc(#loc5796)
    return %9, %5 : tensor<1x1x2048x2048xbf16>, tensor<1x1x2048x2048xbf16> loc(#loc5206)
  } loc(#loc5206)
  func.func private @_pad_1137(%arg0: tensor<1x1x1x2048x2048xbf16> loc(unknown)) -> tensor<1x1x1x2048x2048xbf16> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %1 = stablehlo.pad %arg0, %0, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %2 = "stablehlo.slice"(%1) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %3 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %4 = stablehlo.pad %2, %3, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %5 = "stablehlo.slice"(%4) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %6 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %7 = stablehlo.pad %5, %6, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %8 = "stablehlo.slice"(%7) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %9 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %10 = stablehlo.pad %8, %9, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %11 = "stablehlo.slice"(%10) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %12 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %13 = stablehlo.pad %11, %12, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %14 = "stablehlo.slice"(%13) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %15 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %16 = stablehlo.pad %14, %15, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %17 = "stablehlo.slice"(%16) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %18 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %19 = stablehlo.pad %17, %18, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %20 = "stablehlo.slice"(%19) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %21 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %22 = stablehlo.pad %20, %21, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %23 = "stablehlo.slice"(%22) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %24 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %25 = stablehlo.pad %23, %24, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %26 = "stablehlo.slice"(%25) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    %27 = stablehlo.constant dense<0.000000e+00> : tensor<bf16> loc(#loc5211)
    %28 = stablehlo.pad %26, %27, low = [0, 0, 0, 0, 0], high = [0, 0, 0, 0, 0], interior = [0, 0, 0, 0, 0] : (tensor<1x1x1x2048x2048xbf16>, tensor<bf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6605)
    %29 = "stablehlo.slice"(%28) {limit_indices = dense<[1, 1, 1, 2048, 2048]> : tensor<5xi64>, start_indices = dense<0> : tensor<5xi64>, strides = dense<1> : tensor<5xi64>} : (tensor<1x1x1x2048x2048xbf16>) -> tensor<1x1x1x2048x2048xbf16> loc(#loc6606)
    return %29 : tensor<1x1x1x2048x2048xbf16> loc(#loc5211)
  } loc(#loc5211)
  func.func private @_einsum_1138(%arg0: tensor<1x2048x2048xbf16> loc(unknown), %arg1: tensor<1x2048x2048xbf16> loc(unknown)) -> tensor<2048x2048xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, contracting_dims = [0, 1] x [0, 1], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x2048xbf16>, tensor<1x2048x2048xbf16>) -> tensor<2048x2048xbf16> loc(#loc6607)
    %1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2048x2048xbf16>) -> tensor<2048x2048xbf16> loc(#loc6608)
    return %1 : tensor<2048x2048xbf16> loc(#loc5218)
  } loc(#loc5218)
  func.func private @_einsum_1139(%arg0: tensor<1x2048x51200xbf16> loc(unknown), %arg1: tensor<1x2048x2048xbf16> loc(unknown)) -> tensor<51200x2048xbf16> {
    %0 = stablehlo.dot_general %arg1, %arg0, contracting_dims = [0, 1] x [0, 1], precision = [DEFAULT, DEFAULT] : (tensor<1x2048x2048xbf16>, tensor<1x2048x51200xbf16>) -> tensor<2048x51200xbf16> loc(#loc6609)
    %1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2048x51200xbf16>) -> tensor<51200x2048xbf16> loc(#loc6610)
    return %1 : tensor<51200x2048xbf16> loc(#loc5222)
  } loc(#loc5222)
  func.func private @_where_1140(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1141(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1142(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1143(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : (tensor<i32>) -> tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1144(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6613)
    return %0 : tensor<f32> loc(#loc5237)
  } loc(#loc5237)
  func.func private @isfinite(%arg0: tensor<f32> loc(unknown)) -> tensor<i1> {
    %0 = stablehlo.is_finite %arg0 : (tensor<f32>) -> tensor<i1> loc(#loc6614)
    return %0 : tensor<i1> loc(#loc5250)
  } loc(#loc5250)
  func.func private @nan_to_num(%arg0: tensor<2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<2048xf32> {
    %0 = call @isnan(%arg0) : (tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1145(%0, %1, %arg0) : (tensor<2048xi1>, tensor<f32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6618)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6618)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %7 = call @_where_1146(%5, %6, %2) : (tensor<2048xi1>, tensor<f32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc6617)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6618)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6618)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %12 = call @_where_1147(%10, %11, %7) : (tensor<2048xi1>, tensor<f32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc6617)
    return %12 : tensor<2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan(%arg0: tensor<2048xf32> loc(unknown)) -> tensor<2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6619)
    return %0 : tensor<2048xi1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1145(%arg0: tensor<2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6620)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6621)
    return %1 : tensor<2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1146(%arg0: tensor<2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6620)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6621)
    return %1 : tensor<2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1147(%arg0: tensor<2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6620)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6621)
    return %1 : tensor<2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1148(%arg0: tensor<2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<2048xf32> {
    %0 = call @isnan_1149(%arg0) : (tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1150(%0, %1, %arg0) : (tensor<2048xi1>, tensor<f32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6618)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6618)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %7 = call @_where_1151(%5, %6, %2) : (tensor<2048xi1>, tensor<f32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc6617)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6618)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6618)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %12 = call @_where_1152(%10, %11, %7) : (tensor<2048xi1>, tensor<f32>, tensor<2048xf32>) -> tensor<2048xf32> loc(#loc6617)
    return %12 : tensor<2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1149(%arg0: tensor<2048xf32> loc(unknown)) -> tensor<2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<2048xf32>, tensor<2048xf32>) -> tensor<2048xi1> loc(#loc6619)
    return %0 : tensor<2048xi1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1150(%arg0: tensor<2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6620)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6621)
    return %1 : tensor<2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1151(%arg0: tensor<2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6620)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6621)
    return %1 : tensor<2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1152(%arg0: tensor<2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6620)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6621)
    return %1 : tensor<2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1153(%arg0: tensor<2048x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = call @isnan_1154(%arg0) : (tensor<2048x2048xf32>) -> tensor<2048x2048xi1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1155(%0, %1, %arg0) : (tensor<2048x2048xi1>, tensor<f32>, tensor<2048x2048xf32>) -> tensor<2048x2048xf32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc6618)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<2048x2048xf32>, tensor<2048x2048xf32>) -> tensor<2048x2048xi1> loc(#loc6618)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %7 = call @_where_1156(%5, %6, %2) : (tensor<2048x2048xi1>, tensor<f32>, tensor<2048x2048xf32>) -> tensor<2048x2048xf32> loc(#loc6617)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc6618)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<2048x2048xf32>, tensor<2048x2048xf32>) -> tensor<2048x2048xi1> loc(#loc6618)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %12 = call @_where_1157(%10, %11, %7) : (tensor<2048x2048xi1>, tensor<f32>, tensor<2048x2048xf32>) -> tensor<2048x2048xf32> loc(#loc6617)
    return %12 : tensor<2048x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1154(%arg0: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<2048x2048xf32>, tensor<2048x2048xf32>) -> tensor<2048x2048xi1> loc(#loc6619)
    return %0 : tensor<2048x2048xi1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1155(%arg0: tensor<2048x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc6622)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048x2048xi1>, tensor<2048x2048xf32> loc(#loc6621)
    return %1 : tensor<2048x2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1156(%arg0: tensor<2048x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc6622)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048x2048xi1>, tensor<2048x2048xf32> loc(#loc6621)
    return %1 : tensor<2048x2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1157(%arg0: tensor<2048x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x2048xf32> loc(#loc6622)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048x2048xi1>, tensor<2048x2048xf32> loc(#loc6621)
    return %1 : tensor<2048x2048xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1158(%arg0: tensor<51200xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<51200xf32> {
    %0 = call @isnan_1159(%arg0) : (tensor<51200xf32>) -> tensor<51200xi1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1160(%0, %1, %arg0) : (tensor<51200xi1>, tensor<f32>, tensor<51200xf32>) -> tensor<51200xf32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc6618)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<51200xf32>, tensor<51200xf32>) -> tensor<51200xi1> loc(#loc6618)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %7 = call @_where_1161(%5, %6, %2) : (tensor<51200xi1>, tensor<f32>, tensor<51200xf32>) -> tensor<51200xf32> loc(#loc6617)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc6618)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<51200xf32>, tensor<51200xf32>) -> tensor<51200xi1> loc(#loc6618)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %12 = call @_where_1162(%10, %11, %7) : (tensor<51200xi1>, tensor<f32>, tensor<51200xf32>) -> tensor<51200xf32> loc(#loc6617)
    return %12 : tensor<51200xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1159(%arg0: tensor<51200xf32> loc(unknown)) -> tensor<51200xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<51200xf32>, tensor<51200xf32>) -> tensor<51200xi1> loc(#loc6619)
    return %0 : tensor<51200xi1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1160(%arg0: tensor<51200xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<51200xf32> loc(unknown)) -> tensor<51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc6623)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<51200xi1>, tensor<51200xf32> loc(#loc6621)
    return %1 : tensor<51200xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1161(%arg0: tensor<51200xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<51200xf32> loc(unknown)) -> tensor<51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc6623)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<51200xi1>, tensor<51200xf32> loc(#loc6621)
    return %1 : tensor<51200xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1162(%arg0: tensor<51200xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<51200xf32> loc(unknown)) -> tensor<51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<51200xf32> loc(#loc6623)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<51200xi1>, tensor<51200xf32> loc(#loc6621)
    return %1 : tensor<51200xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1163(%arg0: tensor<2048x51200xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = call @isnan_1164(%arg0) : (tensor<2048x51200xf32>) -> tensor<2048x51200xi1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1165(%0, %1, %arg0) : (tensor<2048x51200xi1>, tensor<f32>, tensor<2048x51200xf32>) -> tensor<2048x51200xf32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc6618)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<2048x51200xf32>, tensor<2048x51200xf32>) -> tensor<2048x51200xi1> loc(#loc6618)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %7 = call @_where_1166(%5, %6, %2) : (tensor<2048x51200xi1>, tensor<f32>, tensor<2048x51200xf32>) -> tensor<2048x51200xf32> loc(#loc6617)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc6618)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<2048x51200xf32>, tensor<2048x51200xf32>) -> tensor<2048x51200xi1> loc(#loc6618)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %12 = call @_where_1167(%10, %11, %7) : (tensor<2048x51200xi1>, tensor<f32>, tensor<2048x51200xf32>) -> tensor<2048x51200xf32> loc(#loc6617)
    return %12 : tensor<2048x51200xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1164(%arg0: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<2048x51200xf32>, tensor<2048x51200xf32>) -> tensor<2048x51200xi1> loc(#loc6619)
    return %0 : tensor<2048x51200xi1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1165(%arg0: tensor<2048x51200xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc6624)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048x51200xi1>, tensor<2048x51200xf32> loc(#loc6621)
    return %1 : tensor<2048x51200xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1166(%arg0: tensor<2048x51200xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc6624)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048x51200xi1>, tensor<2048x51200xf32> loc(#loc6621)
    return %1 : tensor<2048x51200xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1167(%arg0: tensor<2048x51200xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x51200xf32> loc(#loc6624)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<2048x51200xi1>, tensor<2048x51200xf32> loc(#loc6621)
    return %1 : tensor<2048x51200xf32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1168(%arg0: tensor<f32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = call @isnan_1169(%arg0) : (tensor<f32>) -> tensor<i1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1170(%0, %1, %arg0) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.compare  EQ, %2, %3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %5 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %6 = call @_where_1171(%4, %5, %2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %7 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %8 = stablehlo.compare  EQ, %6, %7,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %9 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %10 = call @_where_1172(%8, %9, %6) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    return %10 : tensor<f32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1169(%arg0: tensor<f32> loc(unknown)) -> tensor<i1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6619)
    return %0 : tensor<i1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1170(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1171(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1172(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1173(%arg0: tensor<f32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = call @isnan_1174(%arg0) : (tensor<f32>) -> tensor<i1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1175(%0, %1, %arg0) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.compare  EQ, %2, %3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %5 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %6 = call @_where_1176(%4, %5, %2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %7 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %8 = stablehlo.compare  EQ, %6, %7,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %9 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %10 = call @_where_1177(%8, %9, %6) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    return %10 : tensor<f32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1174(%arg0: tensor<f32> loc(unknown)) -> tensor<i1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6619)
    return %0 : tensor<i1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1175(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1176(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1177(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1178(%arg0: tensor<f32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = call @isnan_1179(%arg0) : (tensor<f32>) -> tensor<i1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1180(%0, %1, %arg0) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.compare  EQ, %2, %3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %5 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %6 = call @_where_1181(%4, %5, %2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %7 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %8 = stablehlo.compare  EQ, %6, %7,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %9 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %10 = call @_where_1182(%8, %9, %6) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    return %10 : tensor<f32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1179(%arg0: tensor<f32> loc(unknown)) -> tensor<i1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6619)
    return %0 : tensor<i1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1180(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1181(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1182(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1183(%arg0: tensor<f32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = call @isnan_1184(%arg0) : (tensor<f32>) -> tensor<i1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1185(%0, %1, %arg0) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.compare  EQ, %2, %3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %5 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %6 = call @_where_1186(%4, %5, %2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %7 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %8 = stablehlo.compare  EQ, %6, %7,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %9 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %10 = call @_where_1187(%8, %9, %6) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    return %10 : tensor<f32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1184(%arg0: tensor<f32> loc(unknown)) -> tensor<i1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6619)
    return %0 : tensor<i1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1185(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1186(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1187(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @nan_to_num_1188(%arg0: tensor<f32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = call @isnan_1189(%arg0) : (tensor<f32>) -> tensor<i1> loc(#loc6615)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6616)
    %2 = call @_where_1190(%0, %1, %arg0) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.compare  EQ, %2, %3,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %5 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6616)
    %6 = call @_where_1191(%4, %5, %2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    %7 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %8 = stablehlo.compare  EQ, %6, %7,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6618)
    %9 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6616)
    %10 = call @_where_1192(%8, %9, %6) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32> loc(#loc6617)
    return %10 : tensor<f32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1189(%arg0: tensor<f32> loc(unknown)) -> tensor<i1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1> loc(#loc6619)
    return %0 : tensor<i1> loc(#loc6615)
  } loc(#loc6615)
  func.func private @_where_1190(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1191(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1192(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6625)
    return %0 : tensor<f32> loc(#loc6617)
  } loc(#loc6617)
  func.func private @_where_1193(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1194(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1195(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1196(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.convert %arg1 : (tensor<i32>) -> tensor<f32> loc(#loc6611)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<i1>, tensor<f32> loc(#loc6612)
    return %1 : tensor<f32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_1197(%arg0: tensor<i1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<f32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<f32> loc(#loc6613)
    return %0 : tensor<f32> loc(#loc5237)
  } loc(#loc5237)
  func.func private @nan_to_num_1198(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1199(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1200(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1201(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1202(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1199(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1200(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1201(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1202(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1203(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1204(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1205(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1206(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1207(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1204(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1205(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1206(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1207(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1208(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1209(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1210(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1211(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1212(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1209(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1210(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1211(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1212(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1213(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1214(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1215(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1216(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1217(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1214(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1215(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1216(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1217(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1218(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1219(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1220(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1221(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1222(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1219(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1220(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1221(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1222(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1223(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1224(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1225(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1226(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1227(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1224(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1225(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1226(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1227(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1228(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1229(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1230(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1231(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1232(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1229(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1230(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1231(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1232(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1233(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1234(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1235(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1236(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1237(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1234(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1235(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1236(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1237(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1238(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1239(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1240(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1241(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1242(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1239(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1240(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1241(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1242(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1243(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1244(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1245(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1246(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1247(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1244(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1245(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1246(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1247(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1248(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1249(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1250(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1251(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1252(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1249(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1250(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1251(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1252(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1253(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1254(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1255(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1256(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1257(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1254(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1255(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1256(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1257(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1258(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1259(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1260(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1261(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1262(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1259(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1260(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1261(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1262(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1263(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1264(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1265(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1266(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1267(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1264(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1265(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1266(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1267(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1268(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1269(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1270(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1271(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1272(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1269(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1270(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1271(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1272(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1273(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1274(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1275(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1276(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1277(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1274(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1275(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1276(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1277(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1278(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1279(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1280(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1281(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1282(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1279(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1280(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1281(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1282(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1283(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1284(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1285(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1286(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1287(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1284(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1285(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1286(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1287(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1288(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1289(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1290(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1291(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1292(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1289(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1290(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1291(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1292(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1293(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1294(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1295(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1296(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1297(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1294(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1295(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1296(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1297(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1298(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1299(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1300(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1301(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1302(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1299(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1300(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1301(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1302(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1303(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1304(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1305(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1306(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1307(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1304(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1305(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1306(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1307(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1308(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1309(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1310(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1311(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1312(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1309(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1310(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1311(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1312(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1313(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1314(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1315(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1316(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1317(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1314(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1315(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1316(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1317(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1318(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1319(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1320(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1321(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1322(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1319(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1320(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1321(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1322(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1323(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1324(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1325(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1326(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1327(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1324(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1325(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1326(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1327(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1328(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1329(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1330(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1331(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1332(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1329(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1330(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1331(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1332(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1333(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1334(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1335(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1336(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1337(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1334(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1335(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1336(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1337(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1338(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1339(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1340(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1341(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1342(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1339(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1340(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1341(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1342(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1343(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1344(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1345(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1346(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1347(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1344(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1345(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1346(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1347(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1348(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1349(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1350(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1351(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1352(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1349(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1350(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1351(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1352(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1353(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1354(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1355(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1356(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1357(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1354(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1355(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1356(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1357(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1358(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1359(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1360(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1361(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1362(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1359(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1360(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1361(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1362(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1363(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1364(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1365(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1366(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1367(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1364(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1365(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1366(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1367(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1368(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1369(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1370(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1371(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1372(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1369(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1370(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1371(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1372(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1373(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1374(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1375(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1376(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1377(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1374(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1375(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1376(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1377(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1378(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1379(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1380(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1381(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1382(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1379(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1380(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1381(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1382(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1383(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1384(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1385(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1386(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1387(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1384(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1385(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1386(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1387(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1388(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1389(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1390(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1391(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1392(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1389(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1390(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1391(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1392(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1393(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1394(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1395(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1396(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1397(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1394(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1395(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1396(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1397(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1398(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1399(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1400(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1401(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1402(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1399(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1400(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1401(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1402(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1403(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1404(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1405(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1406(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1407(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1404(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1405(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1406(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1407(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1408(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1409(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1410(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1411(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1412(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1409(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1410(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1411(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1412(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1413(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1414(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1415(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1416(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1417(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1414(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1415(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1416(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1417(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1418(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1419(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1420(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1421(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1422(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1419(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1420(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1421(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1422(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1423(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1424(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1425(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1426(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1427(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1424(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1425(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1426(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1427(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1428(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1429(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1430(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1431(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1432(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1429(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1430(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1431(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1432(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1433(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1434(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1435(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1436(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1437(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1434(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1435(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1436(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1437(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1438(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1439(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1440(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1441(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1442(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1439(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1440(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1441(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1442(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1443(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1444(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1445(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1446(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1447(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1444(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1445(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1446(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1447(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1448(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1449(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1450(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1451(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1452(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1449(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1450(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1451(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1452(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1453(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1454(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1455(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1456(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1457(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1454(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1455(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1456(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1457(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1458(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1459(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1460(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1461(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1462(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1459(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1460(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1461(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1462(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1463(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1464(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1465(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1466(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1467(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1464(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1465(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1466(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1467(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1468(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1469(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1470(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1471(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1472(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1469(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1470(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1471(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1472(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1473(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1474(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1475(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1476(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1477(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1474(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1475(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1476(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1477(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1478(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1479(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1480(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1481(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1482(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1479(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1480(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1481(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1482(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1483(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1484(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1485(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1486(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1487(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1484(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1485(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1486(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1487(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1488(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1489(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1490(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1491(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1492(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1489(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1490(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1491(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1492(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1493(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1494(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1495(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1496(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1497(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1494(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1495(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1496(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1497(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1498(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1499(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1500(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1501(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1502(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1499(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1500(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1501(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1502(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1503(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1504(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1505(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1506(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1507(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1504(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1505(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1506(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1507(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1508(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1509(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1510(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1511(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1512(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1509(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1510(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1511(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1512(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1513(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1514(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1515(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1516(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1517(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1514(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1515(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1516(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1517(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1518(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1519(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1520(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1521(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1522(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1519(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1520(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1521(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1522(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1523(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1524(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1525(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1526(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1527(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1524(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1525(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1526(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1527(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1528(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1529(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1530(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1531(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1532(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1529(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1530(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1531(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1532(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1533(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1534(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1535(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1536(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1537(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1534(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1535(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1536(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1537(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1538(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1539(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1540(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1541(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1542(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1539(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1540(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1541(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1542(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1543(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1544(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1545(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1546(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1547(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1544(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1545(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1546(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1547(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1548(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1549(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1550(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1551(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1552(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1549(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1550(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1551(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1552(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1553(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1554(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1555(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1556(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1557(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1554(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1555(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1556(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1557(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1558(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1559(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1560(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1561(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1562(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1559(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1560(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1561(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1562(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1563(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1564(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1565(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1566(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1567(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1564(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1565(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1566(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1567(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1568(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1569(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1570(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1571(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1572(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1569(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1570(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1571(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1572(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1573(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1574(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1575(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1576(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1577(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1574(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1575(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1576(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1577(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1578(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1579(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1580(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1581(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1582(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1579(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1580(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1581(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1582(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1583(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1584(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1585(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1586(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1587(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1584(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1585(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1586(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1587(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1588(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1589(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1590(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1591(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1592(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1589(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1590(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1591(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1592(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1593(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1594(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1595(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1596(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1597(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1594(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1595(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1596(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1597(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1598(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1599(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1600(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1601(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1602(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1599(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1600(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1601(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1602(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1603(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1604(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1605(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1606(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1607(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1604(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1605(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1606(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1607(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1608(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1609(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1610(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1611(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1612(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1609(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1610(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1611(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1612(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1613(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1614(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1615(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1616(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1617(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1614(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1615(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1616(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1617(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1618(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1619(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1620(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1621(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1622(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1619(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1620(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1621(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1622(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1623(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1624(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1625(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1626(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1627(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1624(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1625(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1626(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1627(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1628(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1629(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1630(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1631(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1632(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1629(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1630(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1631(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1632(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1633(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1634(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1635(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1636(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1637(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1634(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1635(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1636(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1637(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1638(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1639(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1640(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1641(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1642(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1639(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1640(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1641(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1642(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1643(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1644(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1645(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1646(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1647(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1644(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1645(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1646(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1647(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1648(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1649(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1650(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1651(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1652(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1649(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1650(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1651(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1652(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1653(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1654(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1655(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1656(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1657(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1654(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1655(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1656(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1657(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1658(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1659(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1660(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1661(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1662(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1659(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1660(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1661(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1662(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1663(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1664(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1665(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1666(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1667(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1664(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1665(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1666(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1667(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1668(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1669(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1670(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1671(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1672(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1669(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1670(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1671(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1672(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1673(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1674(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1675(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1676(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1677(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1674(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1675(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1676(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1677(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1678(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1679(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1680(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1681(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1682(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1679(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1680(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1681(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1682(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1683(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1684(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1685(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1686(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1687(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1684(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1685(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1686(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1687(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1688(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1689(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1690(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1691(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1692(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1689(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1690(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1691(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1692(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1693(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1694(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1695(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1696(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1697(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1694(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1695(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1696(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1697(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1698(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1699(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1700(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1701(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1702(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1699(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1700(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1701(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1702(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1703(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1704(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1705(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1706(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1707(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1704(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1705(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1706(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1707(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1708(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1709(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1710(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1711(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1712(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1709(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1710(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1711(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1712(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1713(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1714(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1715(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1716(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1717(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1714(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1715(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1716(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1717(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1718(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1719(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1720(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1721(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1722(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1719(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1720(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1721(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1722(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1723(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1724(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1725(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1726(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1727(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1724(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1725(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1726(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1727(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1728(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1729(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1730(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1731(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1732(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1729(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1730(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1731(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1732(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1733(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1734(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1735(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1736(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1737(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1734(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1735(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1736(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1737(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1738(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1739(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1740(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1741(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1742(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1739(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1740(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1741(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1742(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1743(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1744(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1745(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1746(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1747(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1744(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1745(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1746(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1747(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1748(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1749(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1750(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1751(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1752(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1749(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1750(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1751(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1752(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1753(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1754(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1755(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1756(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1757(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1754(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1755(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1756(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1757(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1758(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1759(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1760(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1761(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1762(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1759(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1760(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1761(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1762(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1763(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1764(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1765(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1766(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1767(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1764(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1765(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1766(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1767(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1768(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1769(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1770(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1771(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1772(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1769(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1770(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1771(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1772(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1773(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1774(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1775(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1776(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1777(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1774(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1775(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1776(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1777(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1778(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1779(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1780(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1781(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1782(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1779(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1780(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1781(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1782(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1783(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1784(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1785(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1786(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1787(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1784(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1785(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1786(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1787(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1788(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1789(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1790(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1791(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1792(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1789(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1790(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1791(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1792(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1793(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1794(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1795(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1796(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1797(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1794(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1795(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1796(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1797(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1798(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1799(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1800(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1801(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1802(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1799(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1800(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1801(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1802(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1803(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1804(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1805(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1806(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1807(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1804(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1805(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1806(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1807(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1808(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1809(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1810(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1811(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1812(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1809(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1810(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1811(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1812(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1813(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1814(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1815(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1816(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1817(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1814(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1815(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1816(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1817(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1818(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1819(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1820(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1821(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1822(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1819(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1820(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1821(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1822(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1823(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1824(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1825(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1826(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1827(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1824(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1825(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1826(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1827(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1828(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1829(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1830(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1831(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1832(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1829(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1830(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1831(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1832(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1833(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1834(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1835(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1836(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1837(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1834(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1835(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1836(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1837(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1838(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1839(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1840(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1841(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1842(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1839(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1840(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1841(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1842(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1843(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1844(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1845(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1846(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1847(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1844(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1845(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1846(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1847(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1848(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1849(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1850(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1851(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1852(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1849(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1850(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1851(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1852(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1853(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1854(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1855(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1856(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1857(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1854(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1855(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1856(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1857(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1858(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1859(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1860(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1861(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1862(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1859(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1860(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1861(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1862(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1863(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1864(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1865(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1866(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1867(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1864(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1865(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1866(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1867(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1868(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1869(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1870(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1871(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1872(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1869(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1870(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1871(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1872(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1873(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1874(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1875(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1876(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1877(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1874(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1875(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1876(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1877(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1878(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1879(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1880(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1881(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1882(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1879(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1880(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1881(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1882(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1883(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1884(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1885(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1886(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1887(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1884(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1885(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1886(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1887(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1888(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1889(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1890(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1891(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1892(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1889(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1890(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1891(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1892(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1893(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1894(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1895(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1896(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1897(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1894(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1895(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1896(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1897(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1898(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1899(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1900(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1901(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1902(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1899(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1900(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1901(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1902(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1903(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1904(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1905(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1906(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1907(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1904(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1905(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1906(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1907(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1908(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1909(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1910(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1911(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1912(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1909(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1910(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1911(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1912(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1913(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1914(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1915(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1916(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1917(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1914(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1915(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1916(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1917(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1918(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1919(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1920(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1921(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1922(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1919(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1920(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1921(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1922(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1923(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1924(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1925(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1926(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1927(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1924(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1925(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1926(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1927(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1928(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1929(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1930(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1931(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1932(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1929(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1930(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1931(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1932(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1933(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1934(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1935(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1936(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1937(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1934(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1935(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1936(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1937(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1938(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1939(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1940(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1941(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1942(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1939(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1940(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1941(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1942(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1943(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1944(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1945(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1946(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1947(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1944(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1945(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1946(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1947(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1948(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1949(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1950(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1951(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1952(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1949(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1950(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1951(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1952(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1953(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_1954(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1955(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1956(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1957(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1954(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1955(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1956(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1957(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1958(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_1959(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1960(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1961(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1962(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1959(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1960(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1961(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1962(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1963(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_1964(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1965(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1966(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1967(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1964(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1965(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1966(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1967(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1968(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_1969(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1970(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1971(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1972(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1969(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1970(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1971(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1972(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1973(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_1974(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1975(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1976(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1977(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1974(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1975(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1976(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1977(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1978(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1979(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1980(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1981(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1982(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1979(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1980(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1981(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1982(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1983(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_1984(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1985(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1986(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1987(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1984(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1985(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1986(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1987(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1988(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1989(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1990(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1991(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1992(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1989(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1990(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1991(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1992(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1993(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1994(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_1995(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_1996(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_1997(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1994(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_1995(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1996(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_1997(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_1998(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_1999(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2000(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2001(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2002(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_1999(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2000(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2001(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2002(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2003(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_2004(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2005(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2006(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2007(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2004(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2005(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2006(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2007(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2008(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_2009(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2010(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2011(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2012(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2009(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2010(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2011(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2012(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2013(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_2014(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2015(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2016(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2017(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2014(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2015(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2016(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2017(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2018(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_2019(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2020(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2021(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2022(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2019(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2020(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2021(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2022(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2023(%arg0: tensor<1x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = call @isnan_2024(%arg0) : (tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2025(%0, %1, %arg0) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2026(%5, %6, %2) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2027(%10, %11, %7) : (tensor<1x8192xi1>, tensor<f32>, tensor<1x8192xf32>) -> tensor<1x8192xf32> loc(#loc6628)
    return %12 : tensor<1x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2024(%arg0: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192xf32>, tensor<1x8192xf32>) -> tensor<1x8192xi1> loc(#loc6630)
    return %0 : tensor<1x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2025(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2026(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2027(%arg0: tensor<1x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192xf32> loc(#loc6631)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<8192xf32>) -> tensor<1x8192xf32> loc(#loc6632)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6633)
    return %2 : tensor<1x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2028(%arg0: tensor<1x2048x8192xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = call @isnan_2029(%arg0) : (tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2030(%0, %1, %arg0) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2031(%5, %6, %2) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x8192xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2032(%10, %11, %7) : (tensor<1x2048x8192xi1>, tensor<f32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6628)
    return %12 : tensor<1x2048x8192xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2029(%arg0: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x8192xf32>, tensor<1x2048x8192xf32>) -> tensor<1x2048x8192xi1> loc(#loc6630)
    return %0 : tensor<1x2048x8192xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2030(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2031(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2032(%arg0: tensor<1x2048x8192xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x8192xf32> loc(#loc6634)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<2048x8192xf32>) -> tensor<1x2048x8192xf32> loc(#loc6635)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6633)
    return %2 : tensor<1x2048x8192xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2033(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_2034(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2035(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2036(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2037(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2034(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2035(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2036(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2037(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2038(%arg0: tensor<1x8192x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = call @isnan_2039(%arg0) : (tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2040(%0, %1, %arg0) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2041(%5, %6, %2) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x8192x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2042(%10, %11, %7) : (tensor<1x8192x2048xi1>, tensor<f32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6628)
    return %12 : tensor<1x8192x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2039(%arg0: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x8192x2048xf32>, tensor<1x8192x2048xf32>) -> tensor<1x8192x2048xi1> loc(#loc6630)
    return %0 : tensor<1x8192x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2040(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2041(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2042(%arg0: tensor<1x8192x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<8192x2048xf32> loc(#loc6638)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2] : (tensor<8192x2048xf32>) -> tensor<1x8192x2048xf32> loc(#loc6639)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6633)
    return %2 : tensor<1x8192x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2043(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_2044(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2045(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2046(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2047(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2044(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2045(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2046(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2047(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2048(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_2049(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2050(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2051(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2052(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2049(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2050(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2051(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2052(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2053(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_2054(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2055(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2056(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2057(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2054(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2055(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2056(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2057(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2058(%arg0: tensor<1x2048xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = call @isnan_2059(%arg0) : (tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2060(%0, %1, %arg0) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2061(%5, %6, %2) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2062(%10, %11, %7) : (tensor<1x2048xi1>, tensor<f32>, tensor<1x2048xf32>) -> tensor<1x2048xf32> loc(#loc6628)
    return %12 : tensor<1x2048xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2059(%arg0: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048xf32>, tensor<1x2048xf32>) -> tensor<1x2048xi1> loc(#loc6630)
    return %0 : tensor<1x2048xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2060(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2061(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2062(%arg0: tensor<1x2048xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048xf32> loc(#loc6636)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<2048xf32>) -> tensor<1x2048xf32> loc(#loc6637)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6633)
    return %2 : tensor<1x2048xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2063(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = call @isnan_2064(%arg0) : (tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2065(%0, %1, %arg0) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2066(%5, %6, %2) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2067(%10, %11, %7) : (tensor<1x3x2048x32x64xi1>, tensor<f32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x3x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2064(%arg0: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x3x2048x32x64xf32>, tensor<1x3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x3x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2065(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2066(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2067(%arg0: tensor<1x3x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<3x2048x32x64xf32> loc(#loc6640)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3, 4] : (tensor<3x2048x32x64xf32>) -> tensor<1x3x2048x32x64xf32> loc(#loc6641)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x3x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2068(%arg0: tensor<1x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = call @isnan_2069(%arg0) : (tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2070(%0, %1, %arg0) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2071(%5, %6, %2) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2072(%10, %11, %7) : (tensor<1x64xi1>, tensor<f32>, tensor<1x64xf32>) -> tensor<1x64xf32> loc(#loc6628)
    return %12 : tensor<1x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2069(%arg0: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xi1> loc(#loc6630)
    return %0 : tensor<1x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2070(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2071(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2072(%arg0: tensor<1x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32> loc(#loc6642)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<64xf32>) -> tensor<1x64xf32> loc(#loc6643)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6633)
    return %2 : tensor<1x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2073(%arg0: tensor<1x2048x32x64xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = call @isnan_2074(%arg0) : (tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2075(%0, %1, %arg0) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2076(%5, %6, %2) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1x2048x32x64xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2077(%10, %11, %7) : (tensor<1x2048x32x64xi1>, tensor<f32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6628)
    return %12 : tensor<1x2048x32x64xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2074(%arg0: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1x2048x32x64xf32>, tensor<1x2048x32x64xf32>) -> tensor<1x2048x32x64xi1> loc(#loc6630)
    return %0 : tensor<1x2048x32x64xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2075(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2076(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2077(%arg0: tensor<1x2048x32x64xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<2048x32x64xf32> loc(#loc6644)
    %1 = stablehlo.broadcast_in_dim %0, dims = [1, 2, 3] : (tensor<2048x32x64xf32>) -> tensor<1x2048x32x64xf32> loc(#loc6645)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6633)
    return %2 : tensor<1x2048x32x64xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2078(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2079(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2080(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2081(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2082(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2079(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2080(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2081(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2082(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2083(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2084(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2085(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2086(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2087(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2084(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2085(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2086(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2087(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2088(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2089(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2090(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2091(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2092(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2089(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2090(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2091(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2092(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2093(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2094(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2095(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2096(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2097(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2094(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2095(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2096(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2097(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2098(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2099(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2100(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2101(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2102(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2099(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2100(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2101(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2102(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2103(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2104(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2105(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2106(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2107(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2104(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2105(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2106(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2107(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2108(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2109(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2110(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2111(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2112(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2109(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2110(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2111(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2112(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2113(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2114(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2115(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2116(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2117(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2114(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2115(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2116(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2117(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2118(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2119(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2120(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2121(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2122(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2119(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2120(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2121(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2122(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2123(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2124(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2125(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2126(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2127(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2124(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2125(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2126(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2127(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2128(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2129(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2130(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2131(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2132(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2129(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2130(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2131(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2132(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2133(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2134(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2135(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2136(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2137(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2134(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2135(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2136(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2137(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2138(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2139(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2140(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2141(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2142(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2139(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2140(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2141(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2142(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2143(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2144(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2145(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2146(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2147(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2144(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2145(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2146(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2147(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2148(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2149(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2150(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2151(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2152(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2149(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2150(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2151(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2152(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2153(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2154(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2155(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2156(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2157(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2154(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2155(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2156(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2157(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2158(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2159(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2160(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2161(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2162(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2159(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2160(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2161(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2162(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2163(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2164(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2165(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2166(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2167(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2164(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2165(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2166(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2167(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2168(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2169(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2170(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2171(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2172(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2169(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2170(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2171(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2172(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2173(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2174(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2175(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2176(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2177(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2174(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2175(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2176(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2177(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2178(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2179(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2180(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2181(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2182(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2179(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2180(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2181(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2182(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2183(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2184(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2185(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2186(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2187(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2184(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2185(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2186(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2187(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2188(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2189(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2190(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2191(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2192(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2189(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2190(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2191(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2192(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2193(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2194(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2195(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2196(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2197(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2194(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2195(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2196(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2197(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2198(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2199(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2200(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2201(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2202(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2199(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2200(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2201(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2202(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2203(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2204(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2205(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2206(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2207(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2204(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2205(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2206(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2207(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2208(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2209(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2210(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2211(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2212(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2209(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2210(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2211(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2212(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2213(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2214(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2215(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2216(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2217(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2214(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2215(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2216(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2217(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2218(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2219(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2220(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2221(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2222(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2219(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2220(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2221(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2222(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2223(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2224(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2225(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2226(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2227(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2224(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2225(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2226(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2227(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2228(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2229(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2230(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2231(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2232(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2229(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2230(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2231(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2232(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2233(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2234(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2235(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2236(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2237(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2234(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2235(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2236(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2237(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2238(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2239(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2240(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2241(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2242(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2239(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2240(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2241(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2242(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2243(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2244(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2245(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2246(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2247(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2244(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2245(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2246(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2247(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2248(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2249(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2250(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2251(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2252(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2249(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2250(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2251(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2252(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2253(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2254(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2255(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2256(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2257(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2254(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2255(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2256(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2257(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2258(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2259(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2260(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2261(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2262(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2259(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2260(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2261(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2262(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2263(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2264(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2265(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2266(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2267(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2264(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2265(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2266(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2267(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2268(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2269(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2270(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2271(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2272(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2269(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2270(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2271(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2272(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2273(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2274(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2275(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2276(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2277(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2274(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2275(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2276(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2277(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2278(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2279(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2280(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2281(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2282(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2279(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2280(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2281(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2282(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2283(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2284(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2285(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2286(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2287(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2284(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2285(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2286(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2287(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2288(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2289(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2290(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2291(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2292(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2289(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2290(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2291(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2292(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2293(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2294(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2295(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2296(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2297(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2294(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2295(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2296(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2297(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2298(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2299(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2300(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2301(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2302(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2299(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2300(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2301(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2302(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2303(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2304(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2305(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2306(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2307(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2304(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2305(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2306(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2307(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2308(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2309(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2310(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2311(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2312(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2309(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2310(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2311(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2312(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2313(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2314(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2315(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2316(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2317(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2314(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2315(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2316(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2317(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2318(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2319(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2320(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2321(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2322(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2319(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2320(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2321(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2322(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2323(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2324(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2325(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2326(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2327(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2324(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2325(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2326(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2327(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2328(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2329(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2330(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2331(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2332(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2329(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2330(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2331(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2332(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2333(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2334(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2335(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2336(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2337(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2334(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2335(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2336(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2337(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2338(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2339(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2340(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2341(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2342(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2339(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2340(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2341(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2342(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2343(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2344(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2345(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2346(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2347(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2344(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2345(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2346(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2347(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2348(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2349(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2350(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2351(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2352(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2349(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2350(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2351(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2352(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2353(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2354(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2355(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2356(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2357(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2354(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2355(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2356(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2357(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2358(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2359(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2360(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2361(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2362(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2359(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2360(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2361(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2362(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2363(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2364(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2365(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2366(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2367(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2364(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2365(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2366(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2367(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2368(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2369(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2370(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2371(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2372(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2369(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2370(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2371(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2372(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2373(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2374(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2375(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2376(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2377(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2374(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2375(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2376(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2377(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2378(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2379(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2380(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2381(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2382(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2379(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2380(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2381(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2382(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2383(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2384(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2385(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2386(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2387(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2384(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2385(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2386(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2387(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2388(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2389(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2390(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2391(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2392(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2389(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2390(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2391(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2392(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2393(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2394(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2395(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2396(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2397(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2394(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2395(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2396(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2397(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2398(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2399(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2400(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2401(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2402(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2399(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2400(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2401(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2402(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2403(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2404(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2405(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2406(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2407(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2404(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2405(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2406(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2407(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2408(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2409(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2410(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2411(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2412(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2409(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2410(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2411(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2412(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2413(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2414(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2415(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2416(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2417(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2414(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2415(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2416(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2417(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2418(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2419(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2420(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2421(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2422(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2419(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2420(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2421(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2422(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2423(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2424(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2425(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2426(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2427(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2424(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2425(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2426(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2427(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2428(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2429(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2430(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2431(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2432(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2429(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2430(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2431(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2432(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2433(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2434(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2435(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2436(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2437(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2434(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2435(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2436(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2437(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2438(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2439(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2440(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2441(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2442(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2439(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2440(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2441(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2442(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2443(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2444(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2445(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2446(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2447(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2444(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2445(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2446(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2447(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2448(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2449(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2450(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2451(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2452(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2449(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2450(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2451(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2452(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2453(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2454(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2455(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2456(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2457(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2454(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2455(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2456(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2457(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2458(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2459(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2460(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2461(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2462(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2459(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2460(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2461(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2462(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2463(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2464(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2465(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2466(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2467(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2464(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2465(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2466(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2467(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2468(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2469(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2470(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2471(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2472(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2469(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2470(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2471(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2472(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2473(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2474(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2475(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2476(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2477(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2474(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2475(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2476(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2477(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2478(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2479(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2480(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2481(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2482(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2479(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2480(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2481(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2482(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2483(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2484(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2485(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2486(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2487(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2484(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2485(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2486(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2487(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2488(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2489(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2490(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2491(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2492(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2489(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2490(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2491(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2492(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2493(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2494(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2495(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2496(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2497(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2494(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2495(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2496(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2497(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2498(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2499(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2500(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2501(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2502(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2499(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2500(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2501(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2502(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2503(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2504(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2505(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2506(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2507(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2504(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2505(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2506(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2507(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2508(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2509(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2510(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2511(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2512(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2509(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2510(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2511(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2512(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2513(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2514(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2515(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2516(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2517(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2514(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2515(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2516(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2517(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2518(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2519(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2520(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2521(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2522(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2519(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2520(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2521(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2522(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2523(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2524(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2525(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2526(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2527(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2524(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2525(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2526(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2527(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2528(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2529(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2530(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2531(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2532(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2529(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2530(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2531(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2532(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2533(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2534(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2535(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2536(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2537(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2534(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2535(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2536(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2537(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2538(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2539(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2540(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2541(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2542(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2539(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2540(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2541(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2542(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2543(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2544(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2545(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2546(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2547(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2544(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2545(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2546(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2547(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2548(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2549(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2550(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2551(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2552(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2549(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2550(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2551(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2552(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2553(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2554(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2555(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2556(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2557(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2554(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2555(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2556(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2557(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2558(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2559(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2560(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2561(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2562(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2559(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2560(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2561(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2562(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2563(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2564(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2565(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2566(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2567(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2564(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2565(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2566(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2567(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2568(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2569(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2570(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2571(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2572(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2569(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2570(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2571(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2572(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2573(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2574(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2575(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2576(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2577(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2574(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2575(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2576(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2577(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2578(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2579(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2580(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2581(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2582(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2579(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2580(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2581(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2582(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2583(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2584(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2585(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2586(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2587(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2584(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2585(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2586(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2587(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2588(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2589(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2590(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2591(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2592(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2589(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2590(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2591(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2592(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2593(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2594(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2595(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2596(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2597(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2594(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2595(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2596(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2597(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2598(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2599(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2600(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2601(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2602(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2599(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2600(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2601(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2602(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2603(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2604(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2605(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2606(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2607(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2604(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2605(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2606(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2607(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2608(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2609(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2610(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2611(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2612(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2609(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2610(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2611(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2612(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2613(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2614(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2615(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2616(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2617(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2614(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2615(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2616(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2617(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2618(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2619(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2620(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2621(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2622(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2619(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2620(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2621(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2622(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2623(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2624(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2625(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2626(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2627(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2624(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2625(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2626(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2627(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2628(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2629(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2630(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2631(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2632(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2629(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2630(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2631(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2632(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2633(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2634(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2635(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2636(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2637(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2634(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2635(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2636(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2637(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2638(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2639(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2640(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2641(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2642(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2639(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2640(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2641(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2642(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2643(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2644(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2645(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2646(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2647(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2644(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2645(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2646(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2647(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2648(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2649(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2650(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2651(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2652(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2649(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2650(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2651(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2652(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2653(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2654(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2655(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2656(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2657(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2654(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2655(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2656(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2657(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2658(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2659(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2660(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2661(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2662(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2659(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2660(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2661(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2662(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2663(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2664(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2665(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2666(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2667(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2664(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2665(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2666(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2667(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2668(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2669(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2670(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2671(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2672(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2669(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2670(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2671(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2672(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2673(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2674(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2675(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2676(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2677(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2674(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2675(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2676(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2677(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2678(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2679(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2680(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2681(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2682(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2679(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2680(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2681(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2682(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2683(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2684(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2685(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2686(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2687(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2684(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2685(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2686(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2687(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2688(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2689(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2690(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2691(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2692(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2689(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2690(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2691(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2692(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2693(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2694(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2695(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2696(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2697(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2694(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2695(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2696(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2697(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2698(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2699(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2700(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2701(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2702(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2699(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2700(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2701(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2702(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2703(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2704(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2705(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2706(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2707(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2704(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2705(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2706(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2707(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2708(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2709(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2710(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2711(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2712(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2709(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2710(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2711(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2712(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2713(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2714(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2715(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2716(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2717(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2714(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2715(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2716(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2717(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2718(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2719(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2720(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2721(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2722(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2719(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2720(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2721(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2722(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2723(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2724(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2725(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2726(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2727(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2724(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2725(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2726(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2727(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2728(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2729(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2730(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2731(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2732(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2729(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2730(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2731(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2732(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2733(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2734(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2735(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2736(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2737(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2734(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2735(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2736(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2737(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2738(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2739(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2740(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2741(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2742(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2739(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2740(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2741(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2742(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2743(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2744(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2745(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2746(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2747(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2744(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2745(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2746(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2747(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2748(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2749(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2750(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2751(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2752(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2749(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2750(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2751(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2752(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2753(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2754(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2755(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2756(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2757(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2754(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2755(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2756(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2757(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2758(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2759(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2760(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2761(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2762(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2759(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2760(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2761(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2762(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2763(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2764(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2765(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2766(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2767(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2764(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2765(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2766(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2767(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2768(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2769(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2770(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2771(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2772(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2769(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2770(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2771(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2772(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2773(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2774(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2775(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2776(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2777(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2774(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2775(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2776(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2777(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2778(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2779(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2780(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2781(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2782(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2779(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2780(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2781(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2782(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2783(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2784(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2785(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2786(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2787(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2784(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2785(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2786(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2787(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2788(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2789(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2790(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2791(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2792(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2789(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2790(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2791(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2792(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2793(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2794(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2795(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2796(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2797(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2794(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2795(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2796(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2797(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2798(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2799(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2800(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2801(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2802(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2799(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2800(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2801(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2802(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2803(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2804(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2805(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2806(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2807(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2804(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2805(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2806(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2807(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2808(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2809(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2810(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2811(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2812(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2809(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2810(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2811(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2812(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2813(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2814(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2815(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2816(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2817(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2814(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2815(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2816(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2817(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2818(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2819(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2820(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2821(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2822(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2819(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2820(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2821(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2822(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2823(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2824(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2825(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2826(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2827(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2824(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2825(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2826(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2827(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2828(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2829(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2830(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2831(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2832(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2829(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2830(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2831(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2832(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2833(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2834(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2835(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2836(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2837(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2834(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2835(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2836(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2837(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2838(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2839(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2840(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2841(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2842(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2839(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2840(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2841(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2842(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2843(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2844(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2845(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2846(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2847(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2844(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2845(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2846(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2847(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2848(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2849(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2850(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2851(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2852(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2849(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2850(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2851(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2852(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2853(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2854(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2855(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2856(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2857(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2854(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2855(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2856(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2857(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2858(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2859(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2860(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2861(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2862(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2859(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2860(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2861(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2862(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2863(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2864(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2865(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2866(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2867(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2864(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2865(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2866(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2867(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2868(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2869(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2870(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2871(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2872(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2869(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2870(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2871(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2872(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2873(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2874(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2875(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2876(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2877(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2874(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2875(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2876(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2877(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2878(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2879(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2880(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2881(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2882(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2879(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2880(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2881(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2882(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2883(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2884(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2885(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2886(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2887(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2884(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2885(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2886(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2887(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2888(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2889(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2890(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2891(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2892(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2889(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2890(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2891(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2892(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2893(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2894(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2895(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2896(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2897(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2894(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2895(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2896(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2897(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2898(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2899(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2900(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2901(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2902(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2899(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2900(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2901(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2902(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2903(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2904(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2905(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2906(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2907(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2904(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2905(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2906(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2907(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2908(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2909(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2910(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2911(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2912(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2909(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2910(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2911(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2912(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2913(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2914(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2915(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2916(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2917(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2914(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2915(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2916(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2917(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2918(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2919(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2920(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2921(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2922(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2919(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2920(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2921(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2922(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2923(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2924(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2925(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2926(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2927(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2924(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2925(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2926(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2927(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2928(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2929(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2930(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2931(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2932(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2929(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2930(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2931(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2932(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2933(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2934(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2935(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2936(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2937(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2934(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2935(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2936(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2937(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2938(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2939(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2940(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2941(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2942(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2939(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2940(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2941(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2942(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2943(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2944(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2945(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2946(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2947(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2944(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2945(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2946(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2947(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2948(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2949(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2950(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2951(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2952(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2949(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2950(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2951(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2952(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @nan_to_num_2953(%arg0: tensor<1xf32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown), %arg3: tensor<f32> loc(unknown)) -> tensor<1xf32> {
    %0 = call @isnan_2954(%arg0) : (tensor<1xf32>) -> tensor<1xi1> loc(#loc6626)
    %1 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6627)
    %2 = call @_where_2955(%0, %1, %arg0) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %3 = stablehlo.constant dense<0x7F800000> : tensor<f32> loc(#loc5257)
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %5 = stablehlo.compare  EQ, %2, %4,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %6 = stablehlo.convert %arg3 : tensor<f32> loc(#loc6627)
    %7 = call @_where_2956(%5, %6, %2) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    %8 = stablehlo.constant dense<0xFF800000> : tensor<f32> loc(#loc5257)
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6629)
    %10 = stablehlo.compare  EQ, %7, %9,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6629)
    %11 = stablehlo.convert %arg2 : tensor<f32> loc(#loc6627)
    %12 = call @_where_2957(%10, %11, %7) : (tensor<1xi1>, tensor<f32>, tensor<1xf32>) -> tensor<1xf32> loc(#loc6628)
    return %12 : tensor<1xf32> loc(#loc5257)
  } loc(#loc5257)
  func.func private @isnan_2954(%arg0: tensor<1xf32> loc(unknown)) -> tensor<1xi1> {
    %0 = stablehlo.compare  NE, %arg0, %arg0,  FLOAT : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> loc(#loc6630)
    return %0 : tensor<1xi1> loc(#loc6626)
  } loc(#loc6626)
  func.func private @_where_2955(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2956(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2957(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6646)
    %1 = stablehlo.select %arg0, %0, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6647)
    return %1 : tensor<1xf32> loc(#loc6628)
  } loc(#loc6628)
  func.func private @_where_2958(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6648)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6649)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6650)
    return %2 : tensor<1xf32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_2959(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6648)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6649)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6650)
    return %2 : tensor<1xf32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_2960(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.convert %arg1 : tensor<f32> loc(#loc6648)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6649)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6650)
    return %2 : tensor<1xf32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_2961(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.convert %arg1 : (tensor<i32>) -> tensor<f32> loc(#loc6648)
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc6649)
    %2 = stablehlo.select %arg0, %1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6650)
    return %2 : tensor<1xf32> loc(#loc5232)
  } loc(#loc5232)
  func.func private @_where_2962(%arg0: tensor<1xi1> loc(unknown), %arg1: tensor<1xf32> loc(unknown), %arg2: tensor<1xf32> loc(unknown)) -> tensor<1xf32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xf32> loc(#loc6651)
    return %0 : tensor<1xf32> loc(#loc5237)
  } loc(#loc5237)
  func.func private @_where_2963(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048xf32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048xi1> loc(#loc6652)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6653)
    return %1 : tensor<2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2964(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048xf32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048xi1> loc(#loc6652)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6653)
    return %1 : tensor<2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2965(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048x2048xf32> loc(unknown), %arg2: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048x2048xi1> loc(#loc6654)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048x2048xi1>, tensor<2048x2048xf32> loc(#loc6653)
    return %1 : tensor<2048x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2966(%arg0: tensor<i1> loc(unknown), %arg1: tensor<51200xf32> loc(unknown), %arg2: tensor<51200xf32> loc(unknown)) -> tensor<51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<51200xi1> loc(#loc6655)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<51200xi1>, tensor<51200xf32> loc(#loc6653)
    return %1 : tensor<51200xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2967(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048x51200xf32> loc(unknown), %arg2: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048x51200xi1> loc(#loc6656)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048x51200xi1>, tensor<2048x51200xf32> loc(#loc6653)
    return %1 : tensor<2048x51200xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2968(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2969(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2970(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2971(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2972(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2973(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2974(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2975(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2976(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2977(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2978(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2979(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2980(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2981(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2982(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2983(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2984(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2985(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2986(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2987(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2988(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2989(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2990(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2991(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2992(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2993(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2994(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2995(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2996(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2997(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2998(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_2999(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3000(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3001(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3002(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3003(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3004(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3005(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3006(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3007(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3008(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3009(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3010(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3011(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3012(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3013(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3014(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3015(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3016(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3017(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3018(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3019(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3020(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3021(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3022(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3023(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3024(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3025(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3026(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3027(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3028(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3029(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3030(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3031(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3032(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3033(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3034(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3035(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3036(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3037(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3038(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3039(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3040(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3041(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3042(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3043(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3044(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3045(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3046(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3047(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3048(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3049(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3050(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3051(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3052(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3053(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3054(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3055(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3056(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3057(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3058(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3059(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3060(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3061(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3062(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3063(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3064(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3065(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3066(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3067(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3068(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3069(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3070(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3071(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3072(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3073(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3074(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3075(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3076(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3077(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3078(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3079(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3080(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3081(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3082(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3083(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3084(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3085(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3086(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3087(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3088(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3089(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3090(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3091(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3092(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3093(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3094(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3095(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3096(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3097(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3098(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3099(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3100(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3101(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3102(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3103(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3104(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3105(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3106(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3107(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3108(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3109(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3110(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3111(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3112(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3113(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3114(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3115(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3116(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3117(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3118(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3119(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3120(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3121(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3122(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3123(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3124(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3125(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3126(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3127(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3128(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3129(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3130(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3131(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3132(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3133(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3134(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3135(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3136(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3137(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3138(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3139(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3140(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3141(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3142(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3143(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5291)
  } loc(#loc5291)
  func.func private @_where_3144(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc6664)
    return %0 : tensor<i32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3145(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc6664)
    return %0 : tensor<i32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3146(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc6664)
    return %0 : tensor<i32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3147(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048xf32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048xi1> loc(#loc6652)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6653)
    return %1 : tensor<2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3148(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048xf32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048xi1> loc(#loc6652)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6653)
    return %1 : tensor<2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3149(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048x2048xf32> loc(unknown), %arg2: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048x2048xi1> loc(#loc6654)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048x2048xi1>, tensor<2048x2048xf32> loc(#loc6653)
    return %1 : tensor<2048x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3150(%arg0: tensor<i1> loc(unknown), %arg1: tensor<51200xf32> loc(unknown), %arg2: tensor<51200xf32> loc(unknown)) -> tensor<51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<51200xi1> loc(#loc6655)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<51200xi1>, tensor<51200xf32> loc(#loc6653)
    return %1 : tensor<51200xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3151(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048x51200xf32> loc(unknown), %arg2: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048x51200xi1> loc(#loc6656)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048x51200xi1>, tensor<2048x51200xf32> loc(#loc6653)
    return %1 : tensor<2048x51200xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3152(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048xf32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048xi1> loc(#loc6652)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6653)
    return %1 : tensor<2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3153(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048xf32> loc(unknown), %arg2: tensor<2048xf32> loc(unknown)) -> tensor<2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048xi1> loc(#loc6652)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048xi1>, tensor<2048xf32> loc(#loc6653)
    return %1 : tensor<2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3154(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048x2048xf32> loc(unknown), %arg2: tensor<2048x2048xf32> loc(unknown)) -> tensor<2048x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048x2048xi1> loc(#loc6654)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048x2048xi1>, tensor<2048x2048xf32> loc(#loc6653)
    return %1 : tensor<2048x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3155(%arg0: tensor<i1> loc(unknown), %arg1: tensor<51200xf32> loc(unknown), %arg2: tensor<51200xf32> loc(unknown)) -> tensor<51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<51200xi1> loc(#loc6655)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<51200xi1>, tensor<51200xf32> loc(#loc6653)
    return %1 : tensor<51200xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3156(%arg0: tensor<i1> loc(unknown), %arg1: tensor<2048x51200xf32> loc(unknown), %arg2: tensor<2048x51200xf32> loc(unknown)) -> tensor<2048x51200xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<2048x51200xi1> loc(#loc6656)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<2048x51200xi1>, tensor<2048x51200xf32> loc(#loc6653)
    return %1 : tensor<2048x51200xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3157(%arg0: tensor<i1> loc(unknown), %arg1: tensor<i32> loc(unknown), %arg2: tensor<i32> loc(unknown)) -> tensor<i32> {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> loc(#loc6664)
    return %0 : tensor<i32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3158(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1xi32> loc(unknown), %arg2: tensor<1xi32> loc(unknown)) -> tensor<1xi32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6665)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1xi1>, tensor<1xi32> loc(#loc6666)
    return %1 : tensor<1xi32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3159(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1xi32> loc(unknown), %arg2: tensor<1xi32> loc(unknown)) -> tensor<1xi32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6665)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1xi1>, tensor<1xi32> loc(#loc6666)
    return %1 : tensor<1xi32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3160(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1xi32> loc(unknown), %arg2: tensor<1xi32> loc(unknown)) -> tensor<1xi32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6665)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1xi1>, tensor<1xi32> loc(#loc6666)
    return %1 : tensor<1xi32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3161(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3162(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3163(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3164(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3165(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3166(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3167(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3168(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3169(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3170(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3171(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3172(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3173(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3174(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3175(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3176(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3177(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3178(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3179(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3180(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3181(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3182(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3183(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3184(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3185(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3186(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3187(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3188(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3189(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3190(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3191(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3192(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3193(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3194(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3195(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3196(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3197(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3198(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3199(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3200(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3201(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3202(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3203(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3204(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3205(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3206(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3207(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3208(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3209(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3210(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3211(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3212(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3213(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3214(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3215(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3216(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3217(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3218(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3219(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3220(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3221(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3222(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3223(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3224(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3225(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3226(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3227(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3228(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3229(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3230(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3231(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3232(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3233(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3234(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3235(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3236(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3237(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3238(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3239(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3240(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3241(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3242(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3243(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3244(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3245(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3246(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3247(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3248(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3249(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3250(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3251(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3252(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3253(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3254(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3255(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3256(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3257(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3258(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3259(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3260(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3261(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3262(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3263(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3264(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3265(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3266(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3267(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3268(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3269(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3270(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3271(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3272(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3273(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3274(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3275(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3276(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3277(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3278(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3279(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3280(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3281(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3282(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3283(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3284(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3285(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3286(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3287(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3288(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3289(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3290(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3291(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3292(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3293(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3294(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3295(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3296(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3297(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3298(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3299(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3300(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3301(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3302(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3303(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3304(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3305(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3306(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3307(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3308(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3309(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3310(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3311(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3312(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3313(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3314(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3315(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3316(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3317(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3318(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3319(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3320(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3321(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3322(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3323(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3324(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3325(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3326(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3327(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3328(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3329(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3330(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3331(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3332(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3333(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3334(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3335(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3336(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3337(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3338(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3339(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3340(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3341(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3342(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3343(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3344(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3345(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3346(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3347(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3348(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3349(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3350(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3351(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3352(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3353(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3354(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3355(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3356(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3357(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3358(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3359(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3360(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3361(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3362(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3363(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3364(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3365(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3366(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3367(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3368(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3369(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3370(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3371(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3372(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3373(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3374(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3375(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3376(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3377(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3378(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3379(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3380(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3381(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3382(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3383(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3384(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3385(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3386(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3387(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3388(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3389(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3390(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3391(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3392(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3393(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3394(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3395(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3396(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3397(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3398(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3399(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3400(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3401(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3402(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3403(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3404(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3405(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3406(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3407(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3408(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3409(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3410(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3411(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3412(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3413(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3414(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3415(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3416(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3417(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3418(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3419(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3420(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3421(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3422(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3423(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3424(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3425(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3426(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3427(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3428(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3429(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3430(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3431(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3432(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3433(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3434(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3435(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3436(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3437(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3438(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3439(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3440(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3441(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3442(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3443(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3444(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3445(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3446(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3447(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3448(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3449(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3450(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3451(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3452(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3453(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3454(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3455(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3456(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3457(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3458(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3459(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3460(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3461(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3462(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3463(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3464(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3465(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3466(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3467(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3468(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3469(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3470(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3471(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3472(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3473(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3474(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3475(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3476(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3477(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3478(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3479(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3480(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3481(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3482(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3483(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3484(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3485(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3486(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3487(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3488(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3489(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3490(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3491(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3492(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3493(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3494(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3495(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3496(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3497(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3498(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3499(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3500(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3501(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3502(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192xf32> loc(unknown), %arg2: tensor<1x8192xf32> loc(unknown)) -> tensor<1x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192xi1> loc(#loc6657)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192xi1>, tensor<1x8192xf32> loc(#loc6653)
    return %1 : tensor<1x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3503(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x8192xf32> loc(unknown), %arg2: tensor<1x2048x8192xf32> loc(unknown)) -> tensor<1x2048x8192xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x8192xi1> loc(#loc6658)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x8192xi1>, tensor<1x2048x8192xf32> loc(#loc6653)
    return %1 : tensor<1x2048x8192xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3504(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3505(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x8192x2048xf32> loc(unknown), %arg2: tensor<1x8192x2048xf32> loc(unknown)) -> tensor<1x8192x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x8192x2048xi1> loc(#loc6660)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x8192x2048xi1>, tensor<1x8192x2048xf32> loc(#loc6653)
    return %1 : tensor<1x8192x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3506(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3507(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3508(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3509(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048xf32> loc(unknown), %arg2: tensor<1x2048xf32> loc(unknown)) -> tensor<1x2048xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048xi1> loc(#loc6659)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048xi1>, tensor<1x2048xf32> loc(#loc6653)
    return %1 : tensor<1x2048xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3510(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x3x2048x32x64xf32> loc(unknown), %arg2: tensor<1x3x2048x32x64xf32> loc(unknown)) -> tensor<1x3x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x3x2048x32x64xi1> loc(#loc6661)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x3x2048x32x64xi1>, tensor<1x3x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x3x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3511(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x64xf32> loc(unknown), %arg2: tensor<1x64xf32> loc(unknown)) -> tensor<1x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x64xi1> loc(#loc6662)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x64xi1>, tensor<1x64xf32> loc(#loc6653)
    return %1 : tensor<1x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3512(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1x2048x32x64xf32> loc(unknown), %arg2: tensor<1x2048x32x64xf32> loc(unknown)) -> tensor<1x2048x32x64xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1x2048x32x64xi1> loc(#loc6663)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1x2048x32x64xi1>, tensor<1x2048x32x64xf32> loc(#loc6653)
    return %1 : tensor<1x2048x32x64xf32> loc(#loc5302)
  } loc(#loc5302)
  func.func private @_where_3513(%arg0: tensor<i1> loc(unknown), %arg1: tensor<1xi32> loc(unknown), %arg2: tensor<1xi32> loc(unknown)) -> tensor<1xi32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i1>) -> tensor<1xi1> loc(#loc6665)
    %1 = stablehlo.select %0, %arg1, %arg2 : tensor<1xi1>, tensor<1xi32> loc(#loc6666)
    return %1 : tensor<1xi32> loc(#loc5302)
  } loc(#loc5302)
} loc(#loc)
#loc1 = loc("/home/silvasean/paxml/praxis/praxis/py_utils.py":485:0)
#loc2 = loc("/home/silvasean/paxml/paxml/paxml/trainer_lib.py":443:0)
#loc3 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":665:0)
#loc4 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":666:0)
#loc5 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":670:0)
#loc6 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":365:0)
#loc7 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":370:0)
#loc8 = loc("/home/silvasean/paxml/praxis/praxis/layers/base_ops.py":43:0)
#loc9 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":378:0)
#loc10 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":1060:0)
#loc11 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":1064:0)
#loc12 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":618:0)
#loc13 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":121:0)
#loc14 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":126:0)
#loc15 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":128:0)
#loc16 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":129:0)
#loc17 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":130:0)
#loc18 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":99:0)
#loc19 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":100:0)
#loc20 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":101:0)
#loc21 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":102:0)
#loc22 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":160:0)
#loc23 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":625:0)
#loc24 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":627:0)
#loc25 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":629:0)
#loc26 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":667:0)
#loc27 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":666:0)
#loc28 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":473:0)
#loc29 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":679:0)
#loc30 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":680:0)
#loc31 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":503:0)
#loc32 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":685:0)
#loc33 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":218:0)
#loc34 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":434:0)
#loc35 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":39:0)
#loc36 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":38:0)
#loc37 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":382:0)
#loc38 = loc("/home/silvasean/paxml/praxis/praxis/py_utils.py":742:0)
#loc39 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":337:0)
#loc40 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":516:0)
#loc41 = loc("/home/silvasean/paxml/paxml.venv/lib/python3.10/site-packages/flax/core/axes_scan.py":149:0)
#loc42 = loc("/home/silvasean/paxml/praxis/praxis/base_layer.py":694:0)
#loc43 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":217:0)
#loc44 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":684:0)
#loc45 = loc("/home/silvasean/paxml/praxis/praxis/layers/normalizations.py":367:0)
#loc46 = loc("/home/silvasean/paxml/praxis/praxis/layers/normalizations.py":368:0)
#loc47 = loc("/home/silvasean/paxml/praxis/praxis/layers/normalizations.py":369:0)
#loc48 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":436:0)
#loc49 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":437:0)
#loc50 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":435:0)
#loc51 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":327:0)
#loc52 = loc("/home/silvasean/paxml/praxis/praxis/layers/normalizations.py":373:0)
#loc53 = loc("/home/silvasean/paxml/praxis/praxis/layers/normalizations.py":375:0)
#loc54 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":850:0)
#loc55 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":383:0)
#loc56 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":1214:0)
#loc57 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":1309:0)
#loc58 = loc("/home/silvasean/paxml/praxis/praxis/layers/attentions.py":1315:0)
#loc59 = loc("/home/silvasean/paxml/praxis/praxis/layers/linears.py":139:0)
#loc60 = loc("/home/silvasean/paxml/praxis/praxis/layers/activations.py":146:0)
#loc61 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformers.py":426:0)
#loc62 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":36:0)
#loc63 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformers.py":454:0)
#loc64 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformers.py":467:0)
#loc65 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":739:0)
#loc66 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":32:0)
#loc67 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformers.py":1301:0)
#loc68 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":43:0)
#loc69 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":40:0)
#loc70 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":41:0)
#loc71 = loc("/home/silvasean/paxml/praxis/praxis/layers/stats.py":42:0)
#loc72 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":814:0)
#loc73 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":549:0)
#loc74 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":917:0)
#loc75 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":916:0)
#loc76 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":945:0)
#loc77 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":946:0)
#loc78 = loc("/home/silvasean/paxml/praxis/praxis/layers/pipeline.py":947:0)
#loc79 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":538:0)
#loc80 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":541:0)
#loc81 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":235:0)
#loc82 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":281:0)
#loc83 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":282:0)
#loc84 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":286:0)
#loc85 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":285:0)
#loc86 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":304:0)
#loc87 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":303:0)
#loc88 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":308:0)
#loc89 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":312:0)
#loc90 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":311:0)
#loc91 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":314:0)
#loc92 = loc("/home/silvasean/paxml/praxis/praxis/layers/embedding_softmax.py":331:0)
#loc93 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":547:0)
#loc94 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":549:0)
#loc95 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformer_models.py":572:0)
#loc96 = loc("/home/silvasean/paxml/praxis/praxis/layers/models.py":126:0)
#loc97 = loc("/home/silvasean/paxml/praxis/praxis/layers/models.py":125:0)
#loc98 = loc("/home/silvasean/paxml/praxis/praxis/layers/models.py":160:0)
#loc99 = loc("/home/silvasean/paxml/paxml.venv/lib/python3.10/site-packages/flax/core/lift.py":1220:0)
#loc100 = loc("/home/silvasean/paxml/praxis/praxis/schedules.py":97:0)
#loc101 = loc("/home/silvasean/paxml/praxis/praxis/schedules.py":100:0)
#loc102 = loc("/home/silvasean/paxml/praxis/praxis/schedules.py":105:0)
#loc103 = loc("/home/silvasean/paxml/praxis/praxis/schedules.py":106:0)
#loc104 = loc("/home/silvasean/paxml/paxml.venv/lib/python3.10/site-packages/optax/_src/schedule.py":399:0)
#loc105 = loc("/home/silvasean/paxml/praxis/praxis/schedules.py":203:0)
#loc106 = loc("/home/silvasean/paxml/praxis/praxis/schedules.py":202:0)
#loc107 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":48:0)
#loc108 = loc("/home/silvasean/paxml/praxis/praxis/layers/transformers.py":70:0)
#loc109 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":50:0)
#loc110 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":257:0)
#loc111 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":293:0)
#loc112 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":263:0)
#loc113 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":267:0)
#loc114 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":263:0)
#loc115 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":285:0)
#loc116 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":286:0)
#loc117 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":294:0)
#loc118 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":295:0)
#loc119 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":618:0)
#loc120 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":1818:0)
#loc121 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":1806:0)
#loc122 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":302:0)
#loc123 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":1802:0)
#loc124 = loc("/home/silvasean/paxml/praxis/praxis/optimizers.py":625:0)
#loc125 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":341:0)
#loc126 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":350:0)
#loc127 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":398:0)
#loc128 = loc("/home/silvasean/paxml/paxml/paxml/learners.py":409:0)
#loc129 = loc("/home/silvasean/paxml/paxml/paxml/train_states.py":58:0)
#loc130 = loc("/home/silvasean/paxml/paxml/paxml/summary_utils.py":163:0)
#loc131 = loc("/home/silvasean/paxml/paxml/paxml/summary_utils.py":165:0)
#loc132 = loc("/home/silvasean/paxml/paxml/paxml/summary_utils.py":164:0)
#loc133 = loc("/home/silvasean/paxml/paxml/paxml/summary_utils.py":180:0)
#loc134 = loc("pjit(_wrapped_step_fn)/jit(main)/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc135 = loc("pjit(_wrapped_step_fn)/jit(main)/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc2))
#loc136 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc3))
#loc137 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/sub"(#loc4))
#loc138 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/reduce_sum[axes=(0, 1)]"(#loc4))
#loc139 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/div"(#loc5))
#loc140 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/transpose[permutation=(1, 0)]"(#loc6))
#loc141 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_one_hot keep_unused=False inline=False]"(#loc7))
#loc142 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc143 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/mul"(#loc9))
#loc144 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc145 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/slice[start_indices=(0, 0) limit_indices=(2048, 2048) strides=(1, 1)]"(#loc10))
#loc146 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_one_hot keep_unused=False inline=False]"(#loc11))
#loc147 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc148 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc149 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/add"(#loc12))
#loc150 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc13))
#loc151 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(0, 2)]"(#loc14))
#loc152 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/ne"(#loc15))
#loc153 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc15))
#loc154 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 2, 3)]"(#loc16))
#loc155 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/mul"(#loc17))
#loc156 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/iota[dtype=int32 shape=(2048,) dimension=0]"(#loc18))
#loc157 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc18))
#loc158 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc18))
#loc159 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/squeeze[dimensions=(0,)]"(#loc18))
#loc160 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(2048, 1, 1, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc18))
#loc161 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/reshape[new_sizes=(2048, 2048) dimensions=None]"(#loc18))
#loc162 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/iota[dtype=int32 shape=(2048,) dimension=0]"(#loc19))
#loc163 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(2048, 1) broadcast_dimensions=(0,)]"(#loc19))
#loc164 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/reshape[new_sizes=(1, 2048, 1, 1) dimensions=None]"(#loc19))
#loc165 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/squeeze[dimensions=(2,)]"(#loc19))
#loc166 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(1, 2048, 2048, 1) broadcast_dimensions=(0, 1, 3)]"(#loc19))
#loc167 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/reshape[new_sizes=(2048, 2048) dimensions=None]"(#loc19))
#loc168 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lt"(#loc20))
#loc169 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc20))
#loc170 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/mul"(#loc20))
#loc171 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(2, 3)]"(#loc21))
#loc172 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/min"(#loc22))
#loc173 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc174 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3}]"(#loc1))
#loc175 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 2048, 2048) dimensions=None]"(#loc23))
#loc176 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/transpose[permutation=(0, 2, 1, 3, 4)]"(#loc24))
#loc177 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 2048, 2048) dimensions=None]"(#loc25))
#loc178 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc23))
#loc179 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/transpose[permutation=(0, 2, 1, 3)]"(#loc24))
#loc180 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 2048) dimensions=None]"(#loc25))
#loc181 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 1, 2048, 2048) dimensions=None]"(#loc23))
#loc182 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/transpose[permutation=(0, 2, 1, 3, 4, 5)]"(#loc24))
#loc183 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 2048, 2048) dimensions=None]"(#loc25))
#loc184 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=(0, 2, 3, 4)]"(#loc26))
#loc185 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_pad keep_unused=False inline=False]"(#loc27))
#loc186 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/pipeline._get_init_loop_state/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc28))
#loc187 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/iota[dtype=int32 shape=(1,) dimension=0]"(#loc29))
#loc188 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=remainder keep_unused=False inline=False]"(#loc30))
#loc189 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_iteration_inputs/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=remainder keep_unused=False inline=False]"(#loc31))
#loc190 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/iota[dtype=int32 shape=(1, 1, 2048, 2048) dimension=0]"(#loc32))
#loc191 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/eq"(#loc32))
#loc192 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._vmap_gather/add"(#loc33))
#loc193 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._vmap_gather/select_n"(#loc33))
#loc194 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._vmap_gather/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=()]"(#loc33))
#loc195 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/iota[dtype=int32 shape=(1,) dimension=0]"(#loc34))
#loc196 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/iota[dtype=int32 shape=(2048,) dimension=0]"(#loc18))
#loc197 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc18))
#loc198 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc18))
#loc199 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/squeeze[dimensions=(0,)]"(#loc18))
#loc200 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(2048, 1, 1, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc18))
#loc201 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(2048, 2048) dimensions=None]"(#loc18))
#loc202 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/iota[dtype=int32 shape=(2048,) dimension=0]"(#loc19))
#loc203 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(2048, 1) broadcast_dimensions=(0,)]"(#loc19))
#loc204 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(1, 2048, 1, 1) dimensions=None]"(#loc19))
#loc205 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/squeeze[dimensions=(2,)]"(#loc19))
#loc206 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 2048, 2048, 1) broadcast_dimensions=(0, 1, 3)]"(#loc19))
#loc207 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(2048, 2048) dimensions=None]"(#loc19))
#loc208 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/lt"(#loc20))
#loc209 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc20))
#loc210 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/mul"(#loc20))
#loc211 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(2, 3)]"(#loc21))
#loc212 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc22))
#loc213 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc214 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc215 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc216 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc217 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc218 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc219 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc220 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc221 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc222 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc223 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc224 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc225 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc226 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc227 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc228 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc229 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc230 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc231 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc232 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc233 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc234 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc235 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc236 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc237 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc238 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc239 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc240 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc241 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc242 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc243 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc244 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc245 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc246 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc247 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc248 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc249 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc250 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc251 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc252 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc253 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc254 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc255 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc256 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc257 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc258 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc259 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc260 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc261 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc262 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc263 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc264 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc265 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc266 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc267 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc268 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc269 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc270 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc271 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc272 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc273 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc274 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc275 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc276 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc277 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc278 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc279 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc280 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc281 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc282 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc283 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc284 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc285 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc286 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc287 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc288 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc289 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc290 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=() out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=() name=softplus keep_unused=False inline=False]"(#loc37))
#loc291 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc292 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc38))
#loc293 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc39))
#loc294 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_new_loop_state/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=_pad keep_unused=False inline=False]"(#loc40))
#loc295 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/empty[dtype=bfloat16]"(#loc41))
#loc296 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc41))
#loc297 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/empty[dtype=float32]"(#loc41))
#loc298 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=()]"(#loc41))
#loc299 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/empty[dtype=int32]"(#loc41))
#loc300 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc41))
#loc301 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/scan[reverse=False length=1 num_consts=489 num_carry=2 linear=(False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False) unroll=1]"(#loc41))
#loc302 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while[cond_nconsts=0 body_nconsts=489]"(#loc41))
#loc303 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/cond/lt"(#loc42))
#loc304 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/sub"(#loc29))
#loc305 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/max"(#loc29))
#loc306 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=remainder keep_unused=False inline=False]"(#loc30))
#loc307 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc308 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/lt"(#loc31))
#loc309 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/add"(#loc33))
#loc310 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/select_n"(#loc33))
#loc311 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=(0,)]"(#loc33))
#loc312 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/concatenate[dimension=1]"(#loc33))
#loc313 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2, 3, 4, 5), collapsed_slice_dims=(), start_index_map=(0, 1, 2, 3, 4)) slice_sizes=(1, 1, 1, 2048, 2048) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc33))
#loc314 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/squeeze[dimensions=(1,)]"(#loc43))
#loc315 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3, 4}]"(#loc1))
#loc316 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3, 4}]"(#loc1))
#loc317 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/min"(#loc22))
#loc318 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc319 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc320 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc321 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc322 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc323 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc324 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc325 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc326 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc327 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc328 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc329 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc330 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc331 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc332 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc333 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc334 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc335 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc336 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc337 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc338 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc339 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc340 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc341 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc342 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc343 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc344 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc345 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc346 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc347 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc348 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc349 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc350 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=remainder keep_unused=False inline=False]"(#loc31))
#loc351 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/lt"(#loc31))
#loc352 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/add"(#loc31))
#loc353 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/select_n"(#loc31))
#loc354 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc31))
#loc355 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 1, 1, 2048, 2048) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc31))
#loc356 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc44))
#loc357 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3}]"(#loc1))
#loc358 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc359 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc360 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc361 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc362 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/div"(#loc45))
#loc363 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/sub"(#loc46))
#loc364 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/mul"(#loc46))
#loc365 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add"(#loc47))
#loc366 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/rsqrt"(#loc47))
#loc367 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/mul"(#loc47))
#loc368 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/le"(#loc48))
#loc369 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/sub"(#loc29))
#loc370 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/lt"(#loc49))
#loc371 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/and"(#loc50))
#loc372 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc51))
#loc373 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add"(#loc52))
#loc374 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc375 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/mul"(#loc52))
#loc376 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc377 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add"(#loc53))
#loc378 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc379 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc380 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc381 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc382 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc383 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc384 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc385 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc386 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc387 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc388 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc389 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc390 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc391 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc392 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc393 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc394 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc395 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc396 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc397 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc398 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/div"(#loc58))
#loc399 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc400 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc401 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc402 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc403 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc404 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/add"(#loc12))
#loc405 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc406 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc407 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc408 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc409 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/div"(#loc45))
#loc410 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/sub"(#loc46))
#loc411 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/mul"(#loc46))
#loc412 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add"(#loc47))
#loc413 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/rsqrt"(#loc47))
#loc414 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/mul"(#loc47))
#loc415 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add"(#loc52))
#loc416 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc417 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/mul"(#loc52))
#loc418 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc419 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add"(#loc53))
#loc420 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc421 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc422 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc423 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc424 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc425 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc426 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc427 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2, 3), collapsed_slice_dims=(), start_index_map=(0, 1, 2)) slice_sizes=(1, 1, 2048) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc33))
#loc428 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._vmap_gather/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2}]"(#loc1))
#loc429 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2}]"(#loc1))
#loc430 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc431 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc432 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/sub"(#loc62))
#loc433 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc434 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/mul"(#loc63))
#loc435 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc436 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc437 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc438 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc439 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc440 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/mul"(#loc64))
#loc441 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/mul"(#loc9))
#loc442 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/add"(#loc12))
#loc443 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc444 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc445 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc446 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc447 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/div"(#loc45))
#loc448 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/sub"(#loc46))
#loc449 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/mul"(#loc46))
#loc450 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add"(#loc47))
#loc451 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/rsqrt"(#loc47))
#loc452 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/mul"(#loc47))
#loc453 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add"(#loc52))
#loc454 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc455 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/mul"(#loc52))
#loc456 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc457 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add"(#loc53))
#loc458 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc459 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc460 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc461 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc462 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc463 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc464 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc465 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc466 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc467 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc468 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc469 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc470 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc471 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc472 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc473 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc474 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc475 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc476 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc477 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc478 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/div"(#loc58))
#loc479 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc480 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc481 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc482 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc483 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc484 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/add"(#loc12))
#loc485 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc486 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc487 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc488 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc489 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/div"(#loc45))
#loc490 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/sub"(#loc46))
#loc491 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/mul"(#loc46))
#loc492 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add"(#loc47))
#loc493 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/rsqrt"(#loc47))
#loc494 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/mul"(#loc47))
#loc495 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add"(#loc52))
#loc496 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc497 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/mul"(#loc52))
#loc498 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc499 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add"(#loc53))
#loc500 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc501 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc502 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc503 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc504 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc505 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc506 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc507 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc508 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc509 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/sub"(#loc62))
#loc510 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc511 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/mul"(#loc63))
#loc512 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc513 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc514 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc515 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc516 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc517 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/mul"(#loc64))
#loc518 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/mul"(#loc9))
#loc519 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/add"(#loc12))
#loc520 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc521 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc522 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc523 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc524 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/div"(#loc45))
#loc525 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/sub"(#loc46))
#loc526 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/mul"(#loc46))
#loc527 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add"(#loc47))
#loc528 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/rsqrt"(#loc47))
#loc529 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/mul"(#loc47))
#loc530 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add"(#loc52))
#loc531 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc532 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/mul"(#loc52))
#loc533 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc534 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add"(#loc53))
#loc535 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc536 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc537 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc538 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc539 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc540 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc541 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc542 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc543 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc544 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc545 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc546 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc547 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc548 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc549 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc550 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc551 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc552 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc553 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc554 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc555 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/div"(#loc58))
#loc556 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc557 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc558 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc559 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc560 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc561 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/add"(#loc12))
#loc562 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc563 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc564 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc565 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc566 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/div"(#loc45))
#loc567 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/sub"(#loc46))
#loc568 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/mul"(#loc46))
#loc569 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add"(#loc47))
#loc570 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/rsqrt"(#loc47))
#loc571 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/mul"(#loc47))
#loc572 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add"(#loc52))
#loc573 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc574 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/mul"(#loc52))
#loc575 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc576 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add"(#loc53))
#loc577 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc578 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc579 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc580 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc581 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc582 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc583 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc584 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc585 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc586 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/sub"(#loc62))
#loc587 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc588 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/mul"(#loc63))
#loc589 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc590 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc591 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc592 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc593 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc594 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/mul"(#loc64))
#loc595 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/mul"(#loc9))
#loc596 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/add"(#loc12))
#loc597 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc598 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc599 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc600 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc601 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/div"(#loc45))
#loc602 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/sub"(#loc46))
#loc603 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/mul"(#loc46))
#loc604 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add"(#loc47))
#loc605 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/rsqrt"(#loc47))
#loc606 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/mul"(#loc47))
#loc607 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add"(#loc52))
#loc608 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc609 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/mul"(#loc52))
#loc610 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc611 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add"(#loc53))
#loc612 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc613 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc614 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc615 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc616 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc617 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc618 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc619 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc620 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc621 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc622 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc623 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc624 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc625 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc626 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc627 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc628 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc629 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc630 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc631 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc632 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/div"(#loc58))
#loc633 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc634 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc635 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc636 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc637 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc638 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/add"(#loc12))
#loc639 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc640 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc641 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc642 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc643 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/div"(#loc45))
#loc644 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/sub"(#loc46))
#loc645 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/mul"(#loc46))
#loc646 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add"(#loc47))
#loc647 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/rsqrt"(#loc47))
#loc648 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/mul"(#loc47))
#loc649 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add"(#loc52))
#loc650 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc651 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/mul"(#loc52))
#loc652 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc653 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add"(#loc53))
#loc654 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc655 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc656 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc657 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc658 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc659 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc660 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc661 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc662 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc663 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/sub"(#loc62))
#loc664 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc665 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/mul"(#loc63))
#loc666 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc667 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc668 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc669 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc670 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc671 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/mul"(#loc64))
#loc672 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/mul"(#loc9))
#loc673 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/add"(#loc12))
#loc674 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc675 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc676 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc677 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc678 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/div"(#loc45))
#loc679 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/sub"(#loc46))
#loc680 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/mul"(#loc46))
#loc681 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add"(#loc47))
#loc682 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/rsqrt"(#loc47))
#loc683 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/mul"(#loc47))
#loc684 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add"(#loc52))
#loc685 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc686 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/mul"(#loc52))
#loc687 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc688 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add"(#loc53))
#loc689 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc690 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc691 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc692 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc693 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc694 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc695 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc696 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc697 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc698 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc699 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc700 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc701 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc702 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc703 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc704 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc705 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc706 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc707 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc708 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc709 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/div"(#loc58))
#loc710 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc711 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc712 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc713 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc714 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc715 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/add"(#loc12))
#loc716 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc717 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc718 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc719 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc720 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/div"(#loc45))
#loc721 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/sub"(#loc46))
#loc722 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/mul"(#loc46))
#loc723 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add"(#loc47))
#loc724 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/rsqrt"(#loc47))
#loc725 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/mul"(#loc47))
#loc726 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add"(#loc52))
#loc727 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc728 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/mul"(#loc52))
#loc729 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc730 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add"(#loc53))
#loc731 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc732 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc733 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc734 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc735 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc736 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc737 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc738 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc739 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc740 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/sub"(#loc62))
#loc741 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc742 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/mul"(#loc63))
#loc743 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc744 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc745 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc746 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc747 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc748 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/mul"(#loc64))
#loc749 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/mul"(#loc9))
#loc750 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/add"(#loc12))
#loc751 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc752 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc753 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc754 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc755 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/div"(#loc45))
#loc756 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/sub"(#loc46))
#loc757 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/mul"(#loc46))
#loc758 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add"(#loc47))
#loc759 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/rsqrt"(#loc47))
#loc760 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/mul"(#loc47))
#loc761 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add"(#loc52))
#loc762 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc763 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/mul"(#loc52))
#loc764 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc765 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add"(#loc53))
#loc766 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc767 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc768 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc769 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc770 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc771 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc772 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc773 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc774 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc775 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc776 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc777 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc778 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc779 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc780 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc781 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc782 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc783 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc784 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc785 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc786 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/div"(#loc58))
#loc787 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc788 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc789 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc790 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc791 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc792 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/add"(#loc12))
#loc793 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc794 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc795 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc796 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc797 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/div"(#loc45))
#loc798 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/sub"(#loc46))
#loc799 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/mul"(#loc46))
#loc800 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add"(#loc47))
#loc801 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/rsqrt"(#loc47))
#loc802 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/mul"(#loc47))
#loc803 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add"(#loc52))
#loc804 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc805 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/mul"(#loc52))
#loc806 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc807 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add"(#loc53))
#loc808 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc809 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc810 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc811 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc812 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc813 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc814 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc815 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc816 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc817 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/sub"(#loc62))
#loc818 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc819 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/mul"(#loc63))
#loc820 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc821 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc822 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc823 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc824 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc825 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/mul"(#loc64))
#loc826 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/mul"(#loc9))
#loc827 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/add"(#loc12))
#loc828 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc829 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc830 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc831 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc832 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/div"(#loc45))
#loc833 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/sub"(#loc46))
#loc834 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/mul"(#loc46))
#loc835 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add"(#loc47))
#loc836 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/rsqrt"(#loc47))
#loc837 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/mul"(#loc47))
#loc838 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add"(#loc52))
#loc839 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc840 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/mul"(#loc52))
#loc841 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc842 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add"(#loc53))
#loc843 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc844 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc845 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc846 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc847 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc848 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc849 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc850 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc851 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc852 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc853 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc854 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc855 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc856 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc857 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc858 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc859 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc860 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc861 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc862 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc863 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/div"(#loc58))
#loc864 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc865 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc866 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc867 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc868 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc869 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/add"(#loc12))
#loc870 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc871 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc872 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc873 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc874 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/div"(#loc45))
#loc875 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/sub"(#loc46))
#loc876 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/mul"(#loc46))
#loc877 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add"(#loc47))
#loc878 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/rsqrt"(#loc47))
#loc879 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/mul"(#loc47))
#loc880 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add"(#loc52))
#loc881 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc882 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/mul"(#loc52))
#loc883 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc884 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add"(#loc53))
#loc885 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc886 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc887 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc888 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc889 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc890 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc891 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc892 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc893 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc894 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/sub"(#loc62))
#loc895 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc896 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/mul"(#loc63))
#loc897 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc898 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc899 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc900 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc901 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc902 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/mul"(#loc64))
#loc903 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/mul"(#loc9))
#loc904 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/add"(#loc12))
#loc905 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc906 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc907 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc908 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc909 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/div"(#loc45))
#loc910 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/sub"(#loc46))
#loc911 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/mul"(#loc46))
#loc912 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add"(#loc47))
#loc913 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/rsqrt"(#loc47))
#loc914 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/mul"(#loc47))
#loc915 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add"(#loc52))
#loc916 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc917 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/mul"(#loc52))
#loc918 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc919 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add"(#loc53))
#loc920 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc921 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc922 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc923 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc924 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc925 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc926 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc927 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc928 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc929 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc930 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc931 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc932 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc933 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc934 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc935 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc936 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc937 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc938 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc939 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc940 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/div"(#loc58))
#loc941 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc942 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc943 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc944 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc945 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc946 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/add"(#loc12))
#loc947 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc948 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc949 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc950 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc951 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/div"(#loc45))
#loc952 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/sub"(#loc46))
#loc953 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/mul"(#loc46))
#loc954 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add"(#loc47))
#loc955 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/rsqrt"(#loc47))
#loc956 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/mul"(#loc47))
#loc957 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add"(#loc52))
#loc958 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc959 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/mul"(#loc52))
#loc960 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc961 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add"(#loc53))
#loc962 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc963 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc964 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc965 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc966 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc967 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc968 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc969 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc970 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc971 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/sub"(#loc62))
#loc972 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc973 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/mul"(#loc63))
#loc974 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc975 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc976 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc977 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc978 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc979 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/mul"(#loc64))
#loc980 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/mul"(#loc9))
#loc981 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/add"(#loc12))
#loc982 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc983 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc984 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc985 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc986 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/div"(#loc45))
#loc987 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/sub"(#loc46))
#loc988 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/mul"(#loc46))
#loc989 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add"(#loc47))
#loc990 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/rsqrt"(#loc47))
#loc991 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/mul"(#loc47))
#loc992 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add"(#loc52))
#loc993 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc994 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/mul"(#loc52))
#loc995 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc996 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add"(#loc53))
#loc997 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc998 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc999 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1000 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1001 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1002 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1003 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1004 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1005 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1006 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1007 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1008 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1009 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1010 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1011 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1012 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1013 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1014 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1015 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1016 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1017 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1018 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1019 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1020 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1021 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1022 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1023 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/add"(#loc12))
#loc1024 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1025 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1026 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1027 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1028 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/div"(#loc45))
#loc1029 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/sub"(#loc46))
#loc1030 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/mul"(#loc46))
#loc1031 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add"(#loc47))
#loc1032 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1033 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/mul"(#loc47))
#loc1034 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add"(#loc52))
#loc1035 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1036 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/mul"(#loc52))
#loc1037 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1038 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add"(#loc53))
#loc1039 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1040 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1041 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1042 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1043 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1044 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1045 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1046 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1047 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1048 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/sub"(#loc62))
#loc1049 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1050 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/mul"(#loc63))
#loc1051 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1052 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1053 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1054 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1055 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1056 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/mul"(#loc64))
#loc1057 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/mul"(#loc9))
#loc1058 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/add"(#loc12))
#loc1059 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1060 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1061 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1062 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1063 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/div"(#loc45))
#loc1064 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/sub"(#loc46))
#loc1065 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/mul"(#loc46))
#loc1066 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add"(#loc47))
#loc1067 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/rsqrt"(#loc47))
#loc1068 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/mul"(#loc47))
#loc1069 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add"(#loc52))
#loc1070 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1071 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/mul"(#loc52))
#loc1072 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1073 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add"(#loc53))
#loc1074 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1075 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1076 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1077 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1078 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1079 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1080 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1081 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1082 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1083 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1084 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1085 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1086 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1087 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1088 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1089 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1090 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1091 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1092 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1093 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1094 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1095 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1096 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1097 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1098 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1099 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1100 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/add"(#loc12))
#loc1101 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1102 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1103 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1104 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1105 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/div"(#loc45))
#loc1106 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/sub"(#loc46))
#loc1107 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/mul"(#loc46))
#loc1108 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add"(#loc47))
#loc1109 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1110 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/mul"(#loc47))
#loc1111 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add"(#loc52))
#loc1112 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1113 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/mul"(#loc52))
#loc1114 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1115 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add"(#loc53))
#loc1116 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1117 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1118 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1119 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1120 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1121 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1122 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1123 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1124 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1125 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/sub"(#loc62))
#loc1126 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1127 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/mul"(#loc63))
#loc1128 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1129 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1130 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1131 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1132 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1133 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/mul"(#loc64))
#loc1134 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/mul"(#loc9))
#loc1135 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/add"(#loc12))
#loc1136 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1137 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1138 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1139 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1140 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/div"(#loc45))
#loc1141 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/sub"(#loc46))
#loc1142 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/mul"(#loc46))
#loc1143 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add"(#loc47))
#loc1144 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/rsqrt"(#loc47))
#loc1145 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/mul"(#loc47))
#loc1146 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add"(#loc52))
#loc1147 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1148 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/mul"(#loc52))
#loc1149 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1150 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add"(#loc53))
#loc1151 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1152 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1153 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1154 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1155 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1156 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1157 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1158 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1159 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1160 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1161 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1162 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1163 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1164 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1165 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1166 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1167 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1168 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1169 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1170 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1171 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1172 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1173 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1174 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1175 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1176 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1177 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/add"(#loc12))
#loc1178 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1179 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1180 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1181 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1182 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/div"(#loc45))
#loc1183 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/sub"(#loc46))
#loc1184 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/mul"(#loc46))
#loc1185 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add"(#loc47))
#loc1186 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1187 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/mul"(#loc47))
#loc1188 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add"(#loc52))
#loc1189 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1190 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/mul"(#loc52))
#loc1191 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1192 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add"(#loc53))
#loc1193 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1194 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1195 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1196 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1197 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1198 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1199 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1200 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1201 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1202 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/sub"(#loc62))
#loc1203 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1204 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/mul"(#loc63))
#loc1205 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1206 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1207 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1208 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1209 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1210 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/mul"(#loc64))
#loc1211 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/mul"(#loc9))
#loc1212 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/add"(#loc12))
#loc1213 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1214 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1215 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1216 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1217 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/div"(#loc45))
#loc1218 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/sub"(#loc46))
#loc1219 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/mul"(#loc46))
#loc1220 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add"(#loc47))
#loc1221 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/rsqrt"(#loc47))
#loc1222 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/mul"(#loc47))
#loc1223 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add"(#loc52))
#loc1224 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1225 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/mul"(#loc52))
#loc1226 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1227 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add"(#loc53))
#loc1228 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1229 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1230 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1231 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1232 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1233 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1234 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1235 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1236 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1237 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1238 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1239 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1240 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1241 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1242 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1243 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1244 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1245 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1246 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1247 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1248 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1249 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1250 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1251 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1252 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1253 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1254 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/add"(#loc12))
#loc1255 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1256 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1257 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1258 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1259 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/div"(#loc45))
#loc1260 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/sub"(#loc46))
#loc1261 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/mul"(#loc46))
#loc1262 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add"(#loc47))
#loc1263 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1264 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/mul"(#loc47))
#loc1265 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add"(#loc52))
#loc1266 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1267 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/mul"(#loc52))
#loc1268 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1269 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add"(#loc53))
#loc1270 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1271 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1272 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1273 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1274 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1275 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1276 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1277 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1278 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1279 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/sub"(#loc62))
#loc1280 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1281 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/mul"(#loc63))
#loc1282 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1283 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1284 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1285 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1286 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1287 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/mul"(#loc64))
#loc1288 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/mul"(#loc9))
#loc1289 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/add"(#loc12))
#loc1290 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1291 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1292 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1293 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1294 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/div"(#loc45))
#loc1295 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/sub"(#loc46))
#loc1296 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/mul"(#loc46))
#loc1297 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add"(#loc47))
#loc1298 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/rsqrt"(#loc47))
#loc1299 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/mul"(#loc47))
#loc1300 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add"(#loc52))
#loc1301 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1302 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/mul"(#loc52))
#loc1303 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1304 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add"(#loc53))
#loc1305 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1306 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1307 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1308 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1309 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1310 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1311 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1312 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1313 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1314 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1315 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1316 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1317 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1318 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1319 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1320 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1321 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1322 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1323 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1324 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1325 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1326 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1327 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1328 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1329 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1330 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1331 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/add"(#loc12))
#loc1332 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1333 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1334 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1335 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1336 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/div"(#loc45))
#loc1337 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/sub"(#loc46))
#loc1338 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/mul"(#loc46))
#loc1339 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add"(#loc47))
#loc1340 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1341 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/mul"(#loc47))
#loc1342 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add"(#loc52))
#loc1343 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1344 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/mul"(#loc52))
#loc1345 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1346 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add"(#loc53))
#loc1347 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1348 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1349 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1350 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1351 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1352 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1353 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1354 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1355 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1356 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/sub"(#loc62))
#loc1357 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1358 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/mul"(#loc63))
#loc1359 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1360 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1361 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1362 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1363 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1364 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/mul"(#loc64))
#loc1365 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/mul"(#loc9))
#loc1366 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/add"(#loc12))
#loc1367 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1368 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1369 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1370 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1371 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/div"(#loc45))
#loc1372 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/sub"(#loc46))
#loc1373 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/mul"(#loc46))
#loc1374 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add"(#loc47))
#loc1375 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/rsqrt"(#loc47))
#loc1376 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/mul"(#loc47))
#loc1377 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add"(#loc52))
#loc1378 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1379 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/mul"(#loc52))
#loc1380 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1381 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add"(#loc53))
#loc1382 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1383 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1384 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1385 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1386 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1387 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1388 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1389 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1390 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1391 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1392 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1393 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1394 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1395 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1396 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1397 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1398 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1399 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1400 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1401 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1402 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1403 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1404 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1405 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1406 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1407 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1408 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/add"(#loc12))
#loc1409 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1410 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1411 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1412 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1413 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/div"(#loc45))
#loc1414 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/sub"(#loc46))
#loc1415 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/mul"(#loc46))
#loc1416 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add"(#loc47))
#loc1417 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1418 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/mul"(#loc47))
#loc1419 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add"(#loc52))
#loc1420 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1421 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/mul"(#loc52))
#loc1422 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1423 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add"(#loc53))
#loc1424 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1425 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1426 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1427 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1428 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1429 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1430 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1431 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1432 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1433 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/sub"(#loc62))
#loc1434 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1435 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/mul"(#loc63))
#loc1436 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1437 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1438 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1439 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1440 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1441 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/mul"(#loc64))
#loc1442 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/mul"(#loc9))
#loc1443 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/add"(#loc12))
#loc1444 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1445 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1446 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1447 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1448 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/div"(#loc45))
#loc1449 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/sub"(#loc46))
#loc1450 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/mul"(#loc46))
#loc1451 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add"(#loc47))
#loc1452 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/rsqrt"(#loc47))
#loc1453 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/mul"(#loc47))
#loc1454 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add"(#loc52))
#loc1455 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1456 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/mul"(#loc52))
#loc1457 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1458 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add"(#loc53))
#loc1459 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1460 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1461 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1462 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1463 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1464 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1465 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1466 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1467 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1468 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1469 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1470 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1471 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1472 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1473 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1474 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1475 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1476 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1477 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1478 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1479 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1480 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1481 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1482 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1483 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1484 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1485 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/add"(#loc12))
#loc1486 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1487 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1488 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1489 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1490 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/div"(#loc45))
#loc1491 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/sub"(#loc46))
#loc1492 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/mul"(#loc46))
#loc1493 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add"(#loc47))
#loc1494 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1495 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/mul"(#loc47))
#loc1496 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add"(#loc52))
#loc1497 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1498 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/mul"(#loc52))
#loc1499 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1500 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add"(#loc53))
#loc1501 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1502 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1503 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1504 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1505 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1506 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1507 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1508 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1509 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1510 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/sub"(#loc62))
#loc1511 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1512 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/mul"(#loc63))
#loc1513 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1514 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1515 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1516 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1517 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1518 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/mul"(#loc64))
#loc1519 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/mul"(#loc9))
#loc1520 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/add"(#loc12))
#loc1521 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1522 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1523 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1524 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1525 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/div"(#loc45))
#loc1526 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/sub"(#loc46))
#loc1527 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/mul"(#loc46))
#loc1528 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add"(#loc47))
#loc1529 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/rsqrt"(#loc47))
#loc1530 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/mul"(#loc47))
#loc1531 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add"(#loc52))
#loc1532 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1533 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/mul"(#loc52))
#loc1534 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1535 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add"(#loc53))
#loc1536 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1537 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1538 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc1539 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1540 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc1541 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc1542 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc1543 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc1544 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1545 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1546 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc1547 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc1548 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc1549 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc1550 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc1551 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc1552 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc1553 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc1554 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc1555 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc1556 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/div"(#loc58))
#loc1557 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc1558 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc1559 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1560 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1561 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1562 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/add"(#loc12))
#loc1563 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1564 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc1565 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc1566 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1567 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/div"(#loc45))
#loc1568 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/sub"(#loc46))
#loc1569 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/mul"(#loc46))
#loc1570 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add"(#loc47))
#loc1571 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/rsqrt"(#loc47))
#loc1572 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/mul"(#loc47))
#loc1573 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add"(#loc52))
#loc1574 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc1575 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/mul"(#loc52))
#loc1576 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1577 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add"(#loc53))
#loc1578 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1579 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1580 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc1581 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc1582 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc1583 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc1584 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc1585 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc1586 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1587 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/sub"(#loc62))
#loc1588 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc1589 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/mul"(#loc63))
#loc1590 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1591 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1592 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc1593 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc1594 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc1595 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/mul"(#loc64))
#loc1596 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/mul"(#loc9))
#loc1597 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/add"(#loc12))
#loc1598 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_new_loop_state/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False, False, False, False, False, False, False) name=_pad keep_unused=False inline=False]"(#loc40))
#loc1599 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_new_loop_state/slice[start_indices=(0, 0, 0, 0) limit_indices=(1, 1, 2048, 2048) strides=None]"(#loc40))
#loc1600 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/add"(#loc65))
#loc1601 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1602 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1603 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1604 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/sub"(#loc62))
#loc1605 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1606 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/mul"(#loc36))
#loc1607 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/abs"(#loc68))
#loc1608 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1609 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc39))
#loc1610 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1611 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/max"(#loc69))
#loc1612 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/div"(#loc5))
#loc1613 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1614 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/sub"(#loc70))
#loc1615 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/mul"(#loc70))
#loc1616 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/sqrt"(#loc71))
#loc1617 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1618 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1619 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1620 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/sub"(#loc62))
#loc1621 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1622 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/mul"(#loc36))
#loc1623 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/abs"(#loc68))
#loc1624 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1625 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1626 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/max"(#loc69))
#loc1627 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/div"(#loc5))
#loc1628 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1629 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/sub"(#loc70))
#loc1630 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/mul"(#loc70))
#loc1631 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/sqrt"(#loc71))
#loc1632 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1633 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1634 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1635 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/sub"(#loc62))
#loc1636 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1637 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/mul"(#loc36))
#loc1638 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/abs"(#loc68))
#loc1639 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1640 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1641 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/max"(#loc69))
#loc1642 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/div"(#loc5))
#loc1643 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1644 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/sub"(#loc70))
#loc1645 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/mul"(#loc70))
#loc1646 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/sqrt"(#loc71))
#loc1647 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1648 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1649 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1650 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/sub"(#loc62))
#loc1651 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1652 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/mul"(#loc36))
#loc1653 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/abs"(#loc68))
#loc1654 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1655 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1656 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/max"(#loc69))
#loc1657 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/div"(#loc5))
#loc1658 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1659 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/sub"(#loc70))
#loc1660 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/mul"(#loc70))
#loc1661 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/sqrt"(#loc71))
#loc1662 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1663 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1664 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1665 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/sub"(#loc62))
#loc1666 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1667 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/mul"(#loc36))
#loc1668 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/abs"(#loc68))
#loc1669 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1670 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1671 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/max"(#loc69))
#loc1672 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/div"(#loc5))
#loc1673 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1674 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/sub"(#loc70))
#loc1675 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/mul"(#loc70))
#loc1676 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/sqrt"(#loc71))
#loc1677 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1678 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1679 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1680 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/sub"(#loc62))
#loc1681 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1682 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/mul"(#loc36))
#loc1683 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/abs"(#loc68))
#loc1684 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1685 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1686 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/max"(#loc69))
#loc1687 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/div"(#loc5))
#loc1688 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1689 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/sub"(#loc70))
#loc1690 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/mul"(#loc70))
#loc1691 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/sqrt"(#loc71))
#loc1692 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1693 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1694 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1695 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/sub"(#loc62))
#loc1696 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1697 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/mul"(#loc36))
#loc1698 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/abs"(#loc68))
#loc1699 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1700 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1701 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/max"(#loc69))
#loc1702 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/div"(#loc5))
#loc1703 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1704 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/sub"(#loc70))
#loc1705 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/mul"(#loc70))
#loc1706 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/sqrt"(#loc71))
#loc1707 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1708 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1709 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1710 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/sub"(#loc62))
#loc1711 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1712 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/mul"(#loc36))
#loc1713 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/abs"(#loc68))
#loc1714 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1715 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1716 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/max"(#loc69))
#loc1717 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/div"(#loc5))
#loc1718 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1719 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/sub"(#loc70))
#loc1720 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/mul"(#loc70))
#loc1721 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/sqrt"(#loc71))
#loc1722 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1723 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1724 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1725 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/sub"(#loc62))
#loc1726 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1727 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/mul"(#loc36))
#loc1728 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/abs"(#loc68))
#loc1729 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1730 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1731 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/max"(#loc69))
#loc1732 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/div"(#loc5))
#loc1733 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1734 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/sub"(#loc70))
#loc1735 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/mul"(#loc70))
#loc1736 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/sqrt"(#loc71))
#loc1737 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1738 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1739 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1740 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/sub"(#loc62))
#loc1741 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1742 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/mul"(#loc36))
#loc1743 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/abs"(#loc68))
#loc1744 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1745 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1746 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/max"(#loc69))
#loc1747 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/div"(#loc5))
#loc1748 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1749 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/sub"(#loc70))
#loc1750 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/mul"(#loc70))
#loc1751 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/sqrt"(#loc71))
#loc1752 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1753 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1754 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1755 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/sub"(#loc62))
#loc1756 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1757 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/mul"(#loc36))
#loc1758 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/abs"(#loc68))
#loc1759 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1760 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1761 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/max"(#loc69))
#loc1762 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/div"(#loc5))
#loc1763 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1764 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/sub"(#loc70))
#loc1765 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/mul"(#loc70))
#loc1766 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/sqrt"(#loc71))
#loc1767 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1768 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1769 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1770 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/sub"(#loc62))
#loc1771 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1772 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/mul"(#loc36))
#loc1773 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/abs"(#loc68))
#loc1774 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1775 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1776 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/max"(#loc69))
#loc1777 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/div"(#loc5))
#loc1778 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1779 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/sub"(#loc70))
#loc1780 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/mul"(#loc70))
#loc1781 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/sqrt"(#loc71))
#loc1782 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1783 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1784 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1785 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/sub"(#loc62))
#loc1786 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1787 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/mul"(#loc36))
#loc1788 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/abs"(#loc68))
#loc1789 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1790 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1791 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/max"(#loc69))
#loc1792 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/div"(#loc5))
#loc1793 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1794 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/sub"(#loc70))
#loc1795 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/mul"(#loc70))
#loc1796 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/sqrt"(#loc71))
#loc1797 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1798 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1799 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1800 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/sub"(#loc62))
#loc1801 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1802 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/mul"(#loc36))
#loc1803 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/abs"(#loc68))
#loc1804 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1805 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1806 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/max"(#loc69))
#loc1807 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/div"(#loc5))
#loc1808 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1809 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/sub"(#loc70))
#loc1810 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/mul"(#loc70))
#loc1811 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/sqrt"(#loc71))
#loc1812 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1813 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1814 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1815 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/sub"(#loc62))
#loc1816 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1817 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/mul"(#loc36))
#loc1818 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/abs"(#loc68))
#loc1819 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1820 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1821 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/max"(#loc69))
#loc1822 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/div"(#loc5))
#loc1823 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1824 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/sub"(#loc70))
#loc1825 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/mul"(#loc70))
#loc1826 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/sqrt"(#loc71))
#loc1827 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc1828 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc1829 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc1830 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/sub"(#loc62))
#loc1831 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc1832 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/mul"(#loc36))
#loc1833 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/abs"(#loc68))
#loc1834 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc1835 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc1836 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/max"(#loc69))
#loc1837 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/div"(#loc5))
#loc1838 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc1839 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/sub"(#loc70))
#loc1840 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/mul"(#loc70))
#loc1841 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/sqrt"(#loc71))
#loc1842 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc41))
#loc1843 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/lt"(#loc42))
#loc1844 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/convert_element_type[new_dtype=int32 weak_type=False]"(#loc42))
#loc1845 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/add"(#loc42))
#loc1846 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/select_n"(#loc41))
#loc1847 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/dynamic_update_slice"(#loc41))
#loc1848 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=(1,)]"(#loc41))
#loc1849 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc41))
#loc1850 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reduce_sum[axes=(0,)]"(#loc72))
#loc1851 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/div"(#loc72))
#loc1852 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/pipeline._unpack_summary/dynamic_slice[slice_sizes=(1,)]"(#loc73))
#loc1853 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/pipeline._unpack_summary/squeeze[dimensions=(0,)]"(#loc73))
#loc1854 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/slice[start_indices=(0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 2048) strides=None]"(#loc74))
#loc1855 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/squeeze[dimensions=(1,)]"(#loc75))
#loc1856 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 2048, 2048) dimensions=None]"(#loc76))
#loc1857 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/transpose[permutation=(0, 2, 1, 3, 4)]"(#loc77))
#loc1858 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 2048, 2048) dimensions=None]"(#loc78))
#loc1859 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1860 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(2,)]"(#loc45))
#loc1861 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc45))
#loc1862 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1863 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/div"(#loc45))
#loc1864 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/sub"(#loc46))
#loc1865 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/mul"(#loc46))
#loc1866 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/add"(#loc47))
#loc1867 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/rsqrt"(#loc47))
#loc1868 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/div"(#loc47))
#loc1869 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/mul"(#loc47))
#loc1870 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/add"(#loc52))
#loc1871 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc52))
#loc1872 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/mul"(#loc52))
#loc1873 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc53))
#loc1874 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/add"(#loc53))
#loc1875 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc79))
#loc1876 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc80))
#loc1877 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1878 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1879 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/bias/broadcast_in_dim[shape=(1, 1, 51200) broadcast_dimensions=(2,)]"(#loc59))
#loc1880 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/bias/add"(#loc59))
#loc1881 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/div"(#loc81))
#loc1882 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/tanh"(#loc81))
#loc1883 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/sub"(#loc81))
#loc1884 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/mul"(#loc81))
#loc1885 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/convert_element_type[new_dtype=float32 weak_type=False]"(#loc82))
#loc1886 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=log_softmax keep_unused=False inline=False]"(#loc83))
#loc1887 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/squeeze[dimensions=(2,)]"(#loc84))
#loc1888 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_one_hot keep_unused=False inline=False]"(#loc85))
#loc1889 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/mul"(#loc86))
#loc1890 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/reduce_sum[axes=(2,)]"(#loc87))
#loc1891 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/neg"(#loc87))
#loc1892 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/argmax[axes=(2,) index_dtype=int32]"(#loc88))
#loc1893 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc89))
#loc1894 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/convert_element_type[new_dtype=float32 weak_type=False]"(#loc89))
#loc1895 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/mul"(#loc89))
#loc1896 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/reduce_sum[axes=(0, 1, 2)]"(#loc90))
#loc1897 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/convert_element_type[new_dtype=float32 weak_type=False]"(#loc91))
#loc1898 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/reduce_sum[axes=(0, 1, 2)]"(#loc91))
#loc1899 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/add"(#loc92))
#loc1900 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/div"(#loc5))
#loc1901 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/convert_element_type[new_dtype=float32 weak_type=False]"(#loc93))
#loc1902 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/mul"(#loc93))
#loc1903 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/reduce_sum[axes=(1,)]"(#loc94))
#loc1904 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/add"(#loc95))
#loc1905 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/eq"(#loc96))
#loc1906 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/convert_element_type[new_dtype=int32 weak_type=False]"(#loc96))
#loc1907 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/mul"(#loc96))
#loc1908 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/reduce_sum[axes=(0, 1)]"(#loc97))
#loc1909 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/max"(#loc96))
#loc1910 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/convert_element_type[new_dtype=float32 weak_type=False]"(#loc97))
#loc1911 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/div"(#loc97))
#loc1912 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_loss/neg"(#loc98))
#loc1913 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/div"(#loc5))
#loc1914 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=()]"(#loc90))
#loc1915 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/mul"(#loc89))
#loc1916 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/reduce_sum[axes=(0, 2)]"(#loc89))
#loc1917 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc89))
#loc1918 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/neg"(#loc87))
#loc1919 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/broadcast_in_dim[shape=(1, 2048, 51200) broadcast_dimensions=(0, 1)]"(#loc87))
#loc1920 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/mul"(#loc86))
#loc1921 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=log_softmax keep_unused=False inline=False]"(#loc83))
#loc1922 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc82))
#loc1923 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/mul"(#loc81))
#loc1924 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/add_any"(#loc81))
#loc1925 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/div"(#loc81))
#loc1926 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/bias/reduce_sum[axes=(1,)]"(#loc59))
#loc1927 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/bias/reshape[new_sizes=(1, 1, 51200) dimensions=None]"(#loc59))
#loc1928 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/bias/reduce_sum[axes=(0, 1)]"(#loc59))
#loc1929 = loc("pjit(_wrapped_step_fn)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc2))
#loc1930 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1931 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc1932 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(1,)]"(#loc53))
#loc1933 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reshape[new_sizes=(1, 1, 2048) dimensions=None]"(#loc53))
#loc1934 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(0, 1)]"(#loc53))
#loc1935 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/mul"(#loc52))
#loc1936 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(1,)]"(#loc52))
#loc1937 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reshape[new_sizes=(1, 1, 2048) dimensions=None]"(#loc52))
#loc1938 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(0, 1)]"(#loc52))
#loc1939 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/mul"(#loc47))
#loc1940 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(2,)]"(#loc47))
#loc1941 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reshape[new_sizes=(1, 2048, 1) dimensions=None]"(#loc47))
#loc1942 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/div"(#loc45))
#loc1943 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc1944 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(0, 2)]"(#loc45))
#loc1945 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc45))
#loc1946 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=(0, 1)]"(#loc45))
#loc1947 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc1948 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/mul"(#loc46))
#loc1949 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/neg"(#loc46))
#loc1950 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reduce_sum[axes=(2,)]"(#loc46))
#loc1951 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/reshape[new_sizes=(1, 2048, 1) dimensions=None]"(#loc46))
#loc1952 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/add_any"(#loc47))
#loc1953 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/final_ln/add_any"(#loc45))
#loc1954 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1955 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 2048, 2048) dimensions=None]"(#loc78))
#loc1956 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/transpose[permutation=(0, 2, 1, 3, 4)]"(#loc77))
#loc1957 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 2048, 2048) dimensions=None]"(#loc76))
#loc1958 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=(0, 2, 3, 4)]"(#loc75))
#loc1959 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/pad[padding_config=((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc74))
#loc1960 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc41))
#loc1961 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=()]"(#loc41))
#loc1962 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc41))
#loc1963 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=()]"(#loc41))
#loc1964 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=()]"(#loc41))
#loc1965 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=()]"(#loc41))
#loc1966 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=()]"(#loc41))
#loc1967 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=()]"(#loc41))
#loc1968 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=()]"(#loc41))
#loc1969 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=()]"(#loc41))
#loc1970 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/scan[reverse=True length=1 num_consts=178 num_carry=178 linear=(False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, False, False) unroll=1]"(#loc41))
#loc1971 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while[cond_nconsts=0 body_nconsts=229]"(#loc41))
#loc1972 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/cond/lt"(#loc42))
#loc1973 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/sub"(#loc41))
#loc1974 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/lt"(#loc42))
#loc1975 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/convert_element_type[new_dtype=int32 weak_type=False]"(#loc42))
#loc1976 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/add"(#loc42))
#loc1977 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/select_n"(#loc41))
#loc1978 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/dynamic_slice[slice_sizes=(1, 1, 1, 2048, 2048)]"(#loc41))
#loc1979 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/squeeze[dimensions=(0,)]"(#loc41))
#loc1980 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/dynamic_slice[slice_sizes=(1, 1)]"(#loc41))
#loc1981 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/dynamic_slice[slice_sizes=(1,)]"(#loc41))
#loc1982 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/iota[dtype=int32 shape=(1,) dimension=0]"(#loc29))
#loc1983 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/sub"(#loc29))
#loc1984 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat2[prevent_cse=False differentiated=True policy=<function save_from_both_policies.<locals>.policy at 0x7fcc99897910>]"(#loc99))
#loc1985 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/max"(#loc29))
#loc1986 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=remainder keep_unused=False inline=False]"(#loc30))
#loc1987 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=remainder keep_unused=False inline=False]"(#loc31))
#loc1988 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/lt"(#loc31))
#loc1989 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/add"(#loc31))
#loc1990 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/select_n"(#loc31))
#loc1991 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc31))
#loc1992 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 1, 1, 2048, 2048) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc31))
#loc1993 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/iota[dtype=int32 shape=(1, 1, 2048, 2048) dimension=0]"(#loc32))
#loc1994 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/eq"(#loc32))
#loc1995 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc44))
#loc1996 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc1997 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/lt"(#loc31))
#loc1998 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/add"(#loc33))
#loc1999 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/select_n"(#loc33))
#loc2000 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=(0,)]"(#loc33))
#loc2001 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=()]"(#loc33))
#loc2002 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/concatenate[dimension=1]"(#loc33))
#loc2003 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2, 3), collapsed_slice_dims=(), start_index_map=(0, 1, 2)) slice_sizes=(1, 1, 2048) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc33))
#loc2004 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/squeeze[dimensions=(1,)]"(#loc43))
#loc2005 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2}]"(#loc1))
#loc2006 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2, 3, 4, 5), collapsed_slice_dims=(), start_index_map=(0, 1, 2, 3, 4)) slice_sizes=(1, 1, 1, 2048, 2048) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc33))
#loc2007 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._vmap_gather/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3, 4}]"(#loc1))
#loc2008 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/iota[dtype=int32 shape=(1,) dimension=0]"(#loc34))
#loc2009 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/le"(#loc48))
#loc2010 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/sub"(#loc29))
#loc2011 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/lt"(#loc49))
#loc2012 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline.get_valid_microbatch_mask/and"(#loc50))
#loc2013 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3}]"(#loc1))
#loc2014 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2}]"(#loc1))
#loc2015 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/pipeline._shard_dim_by_stages/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims={1, 2, 3, 4}]"(#loc1))
#loc2016 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc51))
#loc2017 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_where keep_unused=False inline=False]"(#loc51))
#loc2018 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/iota[dtype=int32 shape=(2048,) dimension=0]"(#loc18))
#loc2019 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc18))
#loc2020 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc18))
#loc2021 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/squeeze[dimensions=(0,)]"(#loc18))
#loc2022 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(2048, 1, 1, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc18))
#loc2023 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(2048, 2048) dimensions=None]"(#loc18))
#loc2024 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/iota[dtype=int32 shape=(2048,) dimension=0]"(#loc19))
#loc2025 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(2048, 1) broadcast_dimensions=(0,)]"(#loc19))
#loc2026 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(1, 2048, 1, 1) dimensions=None]"(#loc19))
#loc2027 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/squeeze[dimensions=(2,)]"(#loc19))
#loc2028 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 2048, 2048, 1) broadcast_dimensions=(0, 1, 3)]"(#loc19))
#loc2029 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/reshape[new_sizes=(2048, 2048) dimensions=None]"(#loc19))
#loc2030 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/lt"(#loc20))
#loc2031 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc20))
#loc2032 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/mul"(#loc20))
#loc2033 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(2, 3)]"(#loc21))
#loc2034 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc22))
#loc2035 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/min"(#loc22))
#loc2036 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2037 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2038 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2039 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/sub"(#loc62))
#loc2040 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2041 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/mul"(#loc36))
#loc2042 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2043 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2044 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2045 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/max"(#loc69))
#loc2046 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/div"(#loc5))
#loc2047 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2048 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/sub"(#loc70))
#loc2049 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/mul"(#loc70))
#loc2050 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/sqrt"(#loc71))
#loc2051 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/div"(#loc71))
#loc2052 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/abs"(#loc68))
#loc2053 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ge"(#loc68))
#loc2054 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2055 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2056 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/eq"(#loc68))
#loc2057 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2058 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2059 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2060 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2061 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2062 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2063 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/div"(#loc45))
#loc2064 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/sub"(#loc46))
#loc2065 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/mul"(#loc46))
#loc2066 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add"(#loc47))
#loc2067 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/rsqrt"(#loc47))
#loc2068 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/div"(#loc47))
#loc2069 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/mul"(#loc47))
#loc2070 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add"(#loc52))
#loc2071 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2072 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/mul"(#loc52))
#loc2073 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2074 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add"(#loc53))
#loc2075 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2076 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2077 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2078 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2079 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2080 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2081 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2082 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2083 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2084 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2085 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2086 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2087 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2088 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2089 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2090 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2091 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2092 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2093 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2094 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2095 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2096 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2097 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2098 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2099 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2100 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2101 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2102 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2103 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2104 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2105 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2106 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2107 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2108 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/add"(#loc12))
#loc2109 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2110 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2111 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2112 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2113 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2114 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/div"(#loc45))
#loc2115 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/sub"(#loc46))
#loc2116 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/mul"(#loc46))
#loc2117 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add"(#loc47))
#loc2118 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2119 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/div"(#loc47))
#loc2120 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/mul"(#loc47))
#loc2121 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add"(#loc52))
#loc2122 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2123 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/mul"(#loc52))
#loc2124 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2125 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add"(#loc53))
#loc2126 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2127 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2128 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2129 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2130 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2131 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2132 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2133 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2134 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2135 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/sub"(#loc62))
#loc2136 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2137 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/mul"(#loc63))
#loc2138 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2139 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2140 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2141 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2142 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2143 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/mul"(#loc64))
#loc2144 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/mul"(#loc9))
#loc2145 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/add"(#loc12))
#loc2146 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2147 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2148 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2149 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/sub"(#loc62))
#loc2150 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2151 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/mul"(#loc36))
#loc2152 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2153 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2154 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2155 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/max"(#loc69))
#loc2156 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/div"(#loc5))
#loc2157 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2158 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/sub"(#loc70))
#loc2159 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/mul"(#loc70))
#loc2160 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/sqrt"(#loc71))
#loc2161 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/div"(#loc71))
#loc2162 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/abs"(#loc68))
#loc2163 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ge"(#loc68))
#loc2164 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2165 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2166 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/eq"(#loc68))
#loc2167 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2168 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2169 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2170 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2171 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2172 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2173 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/div"(#loc45))
#loc2174 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/sub"(#loc46))
#loc2175 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/mul"(#loc46))
#loc2176 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add"(#loc47))
#loc2177 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/rsqrt"(#loc47))
#loc2178 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/div"(#loc47))
#loc2179 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/mul"(#loc47))
#loc2180 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add"(#loc52))
#loc2181 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2182 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/mul"(#loc52))
#loc2183 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2184 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add"(#loc53))
#loc2185 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2186 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2187 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2188 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2189 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2190 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2191 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2192 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2193 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2194 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2195 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2196 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2197 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2198 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2199 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2200 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2201 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2202 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2203 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2204 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2205 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2206 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2207 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2208 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2209 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2210 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2211 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2212 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2213 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2214 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2215 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2216 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2217 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2218 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/add"(#loc12))
#loc2219 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2220 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2221 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2222 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2223 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2224 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/div"(#loc45))
#loc2225 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/sub"(#loc46))
#loc2226 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/mul"(#loc46))
#loc2227 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add"(#loc47))
#loc2228 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2229 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/div"(#loc47))
#loc2230 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/mul"(#loc47))
#loc2231 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add"(#loc52))
#loc2232 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2233 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/mul"(#loc52))
#loc2234 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2235 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add"(#loc53))
#loc2236 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2237 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2238 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2239 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2240 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2241 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2242 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2243 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2244 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2245 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/sub"(#loc62))
#loc2246 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2247 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/mul"(#loc63))
#loc2248 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2249 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2250 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2251 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2252 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2253 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/mul"(#loc64))
#loc2254 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/mul"(#loc9))
#loc2255 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/add"(#loc12))
#loc2256 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2257 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2258 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2259 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/sub"(#loc62))
#loc2260 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2261 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/mul"(#loc36))
#loc2262 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2263 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2264 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2265 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/max"(#loc69))
#loc2266 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/div"(#loc5))
#loc2267 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2268 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/sub"(#loc70))
#loc2269 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/mul"(#loc70))
#loc2270 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/sqrt"(#loc71))
#loc2271 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/div"(#loc71))
#loc2272 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/abs"(#loc68))
#loc2273 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ge"(#loc68))
#loc2274 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2275 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2276 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/eq"(#loc68))
#loc2277 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2278 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2279 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2280 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2281 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2282 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2283 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/div"(#loc45))
#loc2284 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/sub"(#loc46))
#loc2285 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/mul"(#loc46))
#loc2286 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add"(#loc47))
#loc2287 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/rsqrt"(#loc47))
#loc2288 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/div"(#loc47))
#loc2289 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/mul"(#loc47))
#loc2290 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add"(#loc52))
#loc2291 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2292 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/mul"(#loc52))
#loc2293 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2294 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add"(#loc53))
#loc2295 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2296 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2297 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2298 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2299 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2300 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2301 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2302 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2303 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2304 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2305 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2306 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2307 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2308 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2309 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2310 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2311 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2312 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2313 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2314 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2315 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2316 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2317 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2318 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2319 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2320 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2321 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2322 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2323 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2324 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2325 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2326 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2327 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2328 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/add"(#loc12))
#loc2329 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2330 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2331 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2332 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2333 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2334 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/div"(#loc45))
#loc2335 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/sub"(#loc46))
#loc2336 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/mul"(#loc46))
#loc2337 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add"(#loc47))
#loc2338 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2339 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/div"(#loc47))
#loc2340 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/mul"(#loc47))
#loc2341 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add"(#loc52))
#loc2342 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2343 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/mul"(#loc52))
#loc2344 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2345 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add"(#loc53))
#loc2346 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2347 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2348 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2349 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2350 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2351 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2352 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2353 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2354 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2355 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/sub"(#loc62))
#loc2356 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2357 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/mul"(#loc63))
#loc2358 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2359 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2360 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2361 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2362 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2363 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/mul"(#loc64))
#loc2364 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/mul"(#loc9))
#loc2365 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/add"(#loc12))
#loc2366 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2367 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2368 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2369 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/sub"(#loc62))
#loc2370 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2371 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/mul"(#loc36))
#loc2372 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2373 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2374 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2375 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/max"(#loc69))
#loc2376 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/div"(#loc5))
#loc2377 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2378 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/sub"(#loc70))
#loc2379 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/mul"(#loc70))
#loc2380 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/sqrt"(#loc71))
#loc2381 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/div"(#loc71))
#loc2382 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/abs"(#loc68))
#loc2383 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ge"(#loc68))
#loc2384 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2385 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2386 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/eq"(#loc68))
#loc2387 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2388 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2389 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2390 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2391 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2392 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2393 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/div"(#loc45))
#loc2394 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/sub"(#loc46))
#loc2395 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/mul"(#loc46))
#loc2396 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add"(#loc47))
#loc2397 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/rsqrt"(#loc47))
#loc2398 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/div"(#loc47))
#loc2399 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/mul"(#loc47))
#loc2400 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add"(#loc52))
#loc2401 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2402 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/mul"(#loc52))
#loc2403 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2404 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add"(#loc53))
#loc2405 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2406 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2407 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2408 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2409 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2410 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2411 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2412 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2413 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2414 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2415 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2416 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2417 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2418 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2419 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2420 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2421 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2422 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2423 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2424 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2425 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2426 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2427 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2428 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2429 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2430 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2431 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2432 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2433 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2434 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2435 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2436 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2437 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2438 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/add"(#loc12))
#loc2439 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2440 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2441 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2442 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2443 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2444 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/div"(#loc45))
#loc2445 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/sub"(#loc46))
#loc2446 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/mul"(#loc46))
#loc2447 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add"(#loc47))
#loc2448 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2449 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/div"(#loc47))
#loc2450 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/mul"(#loc47))
#loc2451 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add"(#loc52))
#loc2452 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2453 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/mul"(#loc52))
#loc2454 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2455 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add"(#loc53))
#loc2456 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2457 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2458 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2459 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2460 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2461 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2462 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2463 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2464 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2465 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/sub"(#loc62))
#loc2466 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2467 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/mul"(#loc63))
#loc2468 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2469 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2470 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2471 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2472 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2473 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/mul"(#loc64))
#loc2474 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/mul"(#loc9))
#loc2475 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/add"(#loc12))
#loc2476 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2477 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2478 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2479 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/sub"(#loc62))
#loc2480 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2481 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/mul"(#loc36))
#loc2482 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2483 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2484 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2485 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/max"(#loc69))
#loc2486 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/div"(#loc5))
#loc2487 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2488 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/sub"(#loc70))
#loc2489 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/mul"(#loc70))
#loc2490 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/sqrt"(#loc71))
#loc2491 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/div"(#loc71))
#loc2492 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/abs"(#loc68))
#loc2493 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ge"(#loc68))
#loc2494 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2495 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2496 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/eq"(#loc68))
#loc2497 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2498 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2499 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2500 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2501 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2502 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2503 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/div"(#loc45))
#loc2504 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/sub"(#loc46))
#loc2505 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/mul"(#loc46))
#loc2506 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add"(#loc47))
#loc2507 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/rsqrt"(#loc47))
#loc2508 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/div"(#loc47))
#loc2509 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/mul"(#loc47))
#loc2510 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add"(#loc52))
#loc2511 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2512 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/mul"(#loc52))
#loc2513 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2514 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add"(#loc53))
#loc2515 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2516 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2517 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2518 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2519 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2520 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2521 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2522 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2523 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2524 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2525 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2526 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2527 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2528 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2529 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2530 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2531 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2532 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2533 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2534 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2535 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2536 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2537 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2538 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2539 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2540 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2541 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2542 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2543 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2544 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2545 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2546 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2547 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2548 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/add"(#loc12))
#loc2549 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2550 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2551 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2552 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2553 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2554 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/div"(#loc45))
#loc2555 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/sub"(#loc46))
#loc2556 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/mul"(#loc46))
#loc2557 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add"(#loc47))
#loc2558 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2559 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/div"(#loc47))
#loc2560 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/mul"(#loc47))
#loc2561 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add"(#loc52))
#loc2562 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2563 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/mul"(#loc52))
#loc2564 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2565 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add"(#loc53))
#loc2566 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2567 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2568 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2569 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2570 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2571 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2572 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2573 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2574 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2575 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/sub"(#loc62))
#loc2576 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2577 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/mul"(#loc63))
#loc2578 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2579 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2580 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2581 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2582 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2583 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/mul"(#loc64))
#loc2584 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/mul"(#loc9))
#loc2585 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/add"(#loc12))
#loc2586 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2587 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2588 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2589 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/sub"(#loc62))
#loc2590 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2591 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/mul"(#loc36))
#loc2592 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2593 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2594 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2595 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/max"(#loc69))
#loc2596 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/div"(#loc5))
#loc2597 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2598 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/sub"(#loc70))
#loc2599 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/mul"(#loc70))
#loc2600 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/sqrt"(#loc71))
#loc2601 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/div"(#loc71))
#loc2602 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/abs"(#loc68))
#loc2603 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ge"(#loc68))
#loc2604 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2605 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2606 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/eq"(#loc68))
#loc2607 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2608 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2609 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2610 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2611 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2612 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2613 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/div"(#loc45))
#loc2614 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/sub"(#loc46))
#loc2615 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/mul"(#loc46))
#loc2616 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add"(#loc47))
#loc2617 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/rsqrt"(#loc47))
#loc2618 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/div"(#loc47))
#loc2619 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/mul"(#loc47))
#loc2620 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add"(#loc52))
#loc2621 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2622 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/mul"(#loc52))
#loc2623 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2624 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add"(#loc53))
#loc2625 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2626 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2627 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2628 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2629 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2630 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2631 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2632 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2633 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2634 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2635 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2636 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2637 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2638 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2639 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2640 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2641 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2642 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2643 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2644 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2645 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2646 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2647 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2648 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2649 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2650 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2651 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2652 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2653 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2654 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2655 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2656 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2657 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2658 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/add"(#loc12))
#loc2659 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2660 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2661 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2662 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2663 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2664 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/div"(#loc45))
#loc2665 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/sub"(#loc46))
#loc2666 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/mul"(#loc46))
#loc2667 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add"(#loc47))
#loc2668 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2669 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/div"(#loc47))
#loc2670 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/mul"(#loc47))
#loc2671 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add"(#loc52))
#loc2672 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2673 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/mul"(#loc52))
#loc2674 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2675 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add"(#loc53))
#loc2676 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2677 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2678 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2679 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2680 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2681 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2682 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2683 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2684 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2685 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/sub"(#loc62))
#loc2686 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2687 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/mul"(#loc63))
#loc2688 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2689 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2690 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2691 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2692 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2693 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/mul"(#loc64))
#loc2694 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/mul"(#loc9))
#loc2695 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/add"(#loc12))
#loc2696 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2697 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2698 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2699 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/sub"(#loc62))
#loc2700 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2701 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/mul"(#loc36))
#loc2702 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2703 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2704 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2705 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/max"(#loc69))
#loc2706 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/div"(#loc5))
#loc2707 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2708 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/sub"(#loc70))
#loc2709 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/mul"(#loc70))
#loc2710 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/sqrt"(#loc71))
#loc2711 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/div"(#loc71))
#loc2712 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/abs"(#loc68))
#loc2713 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ge"(#loc68))
#loc2714 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2715 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2716 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/eq"(#loc68))
#loc2717 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2718 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2719 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2720 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2721 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2722 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2723 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/div"(#loc45))
#loc2724 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/sub"(#loc46))
#loc2725 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/mul"(#loc46))
#loc2726 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add"(#loc47))
#loc2727 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/rsqrt"(#loc47))
#loc2728 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/div"(#loc47))
#loc2729 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/mul"(#loc47))
#loc2730 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add"(#loc52))
#loc2731 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2732 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/mul"(#loc52))
#loc2733 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2734 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add"(#loc53))
#loc2735 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2736 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2737 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2738 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2739 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2740 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2741 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2742 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2743 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2744 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2745 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2746 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2747 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2748 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2749 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2750 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2751 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2752 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2753 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2754 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2755 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2756 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2757 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2758 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2759 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2760 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2761 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2762 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2763 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2764 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2765 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2766 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2767 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2768 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/add"(#loc12))
#loc2769 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2770 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2771 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2772 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2773 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2774 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/div"(#loc45))
#loc2775 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/sub"(#loc46))
#loc2776 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/mul"(#loc46))
#loc2777 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add"(#loc47))
#loc2778 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2779 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/div"(#loc47))
#loc2780 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/mul"(#loc47))
#loc2781 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add"(#loc52))
#loc2782 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2783 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/mul"(#loc52))
#loc2784 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2785 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add"(#loc53))
#loc2786 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2787 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2788 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2789 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2790 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2791 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2792 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2793 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2794 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2795 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/sub"(#loc62))
#loc2796 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2797 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/mul"(#loc63))
#loc2798 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2799 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2800 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2801 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2802 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2803 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/mul"(#loc64))
#loc2804 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/mul"(#loc9))
#loc2805 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/add"(#loc12))
#loc2806 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2807 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2808 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2809 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/sub"(#loc62))
#loc2810 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2811 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/mul"(#loc36))
#loc2812 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2813 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2814 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2815 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/max"(#loc69))
#loc2816 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/div"(#loc5))
#loc2817 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2818 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/sub"(#loc70))
#loc2819 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/mul"(#loc70))
#loc2820 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/sqrt"(#loc71))
#loc2821 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/div"(#loc71))
#loc2822 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/abs"(#loc68))
#loc2823 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ge"(#loc68))
#loc2824 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2825 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2826 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/eq"(#loc68))
#loc2827 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2828 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2829 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2830 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2831 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2832 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2833 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/div"(#loc45))
#loc2834 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/sub"(#loc46))
#loc2835 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/mul"(#loc46))
#loc2836 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add"(#loc47))
#loc2837 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/rsqrt"(#loc47))
#loc2838 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/div"(#loc47))
#loc2839 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/mul"(#loc47))
#loc2840 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add"(#loc52))
#loc2841 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2842 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/mul"(#loc52))
#loc2843 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2844 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add"(#loc53))
#loc2845 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2846 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2847 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2848 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2849 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2850 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2851 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2852 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2853 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2854 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2855 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2856 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2857 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2858 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2859 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2860 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2861 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2862 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2863 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2864 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2865 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2866 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2867 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2868 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2869 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2870 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2871 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2872 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2873 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2874 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2875 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2876 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2877 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2878 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/add"(#loc12))
#loc2879 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2880 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2881 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2882 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2883 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2884 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/div"(#loc45))
#loc2885 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/sub"(#loc46))
#loc2886 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/mul"(#loc46))
#loc2887 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add"(#loc47))
#loc2888 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2889 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/div"(#loc47))
#loc2890 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/mul"(#loc47))
#loc2891 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add"(#loc52))
#loc2892 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2893 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/mul"(#loc52))
#loc2894 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2895 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add"(#loc53))
#loc2896 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2897 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2898 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc2899 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc2900 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc2901 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc2902 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc2903 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc2904 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2905 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/sub"(#loc62))
#loc2906 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc2907 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/mul"(#loc63))
#loc2908 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2909 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2910 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2911 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc2912 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc2913 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/mul"(#loc64))
#loc2914 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/mul"(#loc9))
#loc2915 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/add"(#loc12))
#loc2916 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc2917 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc2918 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc2919 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/sub"(#loc62))
#loc2920 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc2921 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/mul"(#loc36))
#loc2922 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc2923 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc2924 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc2925 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/max"(#loc69))
#loc2926 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/div"(#loc5))
#loc2927 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc2928 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/sub"(#loc70))
#loc2929 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/mul"(#loc70))
#loc2930 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/sqrt"(#loc71))
#loc2931 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/div"(#loc71))
#loc2932 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/abs"(#loc68))
#loc2933 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ge"(#loc68))
#loc2934 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc2935 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc2936 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/eq"(#loc68))
#loc2937 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc2938 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc2939 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2940 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2941 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2942 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2943 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/div"(#loc45))
#loc2944 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/sub"(#loc46))
#loc2945 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/mul"(#loc46))
#loc2946 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add"(#loc47))
#loc2947 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/rsqrt"(#loc47))
#loc2948 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/div"(#loc47))
#loc2949 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/mul"(#loc47))
#loc2950 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add"(#loc52))
#loc2951 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc2952 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/mul"(#loc52))
#loc2953 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc2954 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add"(#loc53))
#loc2955 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2956 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2957 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc2958 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2959 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc2960 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2961 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc2962 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc2963 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc2964 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc2965 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2966 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc2967 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc2968 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc2969 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc2970 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc2971 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc2972 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc2973 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc2974 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc2975 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc2976 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc2977 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc2978 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc2979 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc2980 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/div"(#loc58))
#loc2981 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc2982 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc2983 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc2984 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc2985 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2986 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc2987 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc2988 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/add"(#loc12))
#loc2989 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc2990 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc2991 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc2992 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc2993 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc2994 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/div"(#loc45))
#loc2995 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/sub"(#loc46))
#loc2996 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/mul"(#loc46))
#loc2997 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add"(#loc47))
#loc2998 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/rsqrt"(#loc47))
#loc2999 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/div"(#loc47))
#loc3000 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/mul"(#loc47))
#loc3001 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add"(#loc52))
#loc3002 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3003 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/mul"(#loc52))
#loc3004 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3005 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add"(#loc53))
#loc3006 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3007 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3008 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3009 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3010 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3011 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3012 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3013 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3014 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3015 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/sub"(#loc62))
#loc3016 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3017 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/mul"(#loc63))
#loc3018 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3019 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3020 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3021 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3022 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3023 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/mul"(#loc64))
#loc3024 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/mul"(#loc9))
#loc3025 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/add"(#loc12))
#loc3026 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3027 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3028 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3029 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/sub"(#loc62))
#loc3030 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3031 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/mul"(#loc36))
#loc3032 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3033 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3034 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3035 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/max"(#loc69))
#loc3036 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/div"(#loc5))
#loc3037 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3038 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/sub"(#loc70))
#loc3039 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/mul"(#loc70))
#loc3040 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/sqrt"(#loc71))
#loc3041 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/div"(#loc71))
#loc3042 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/abs"(#loc68))
#loc3043 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ge"(#loc68))
#loc3044 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3045 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3046 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/eq"(#loc68))
#loc3047 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3048 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3049 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3050 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3051 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3052 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3053 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/div"(#loc45))
#loc3054 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/sub"(#loc46))
#loc3055 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/mul"(#loc46))
#loc3056 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add"(#loc47))
#loc3057 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/rsqrt"(#loc47))
#loc3058 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/div"(#loc47))
#loc3059 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/mul"(#loc47))
#loc3060 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add"(#loc52))
#loc3061 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3062 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/mul"(#loc52))
#loc3063 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3064 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add"(#loc53))
#loc3065 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3066 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3067 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3068 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3069 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3070 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3071 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3072 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3073 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3074 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3075 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3076 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3077 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3078 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3079 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3080 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3081 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3082 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3083 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3084 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3085 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3086 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3087 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3088 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3089 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3090 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3091 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3092 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3093 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3094 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3095 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3096 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3097 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3098 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/add"(#loc12))
#loc3099 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3100 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3101 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3102 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3103 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3104 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/div"(#loc45))
#loc3105 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/sub"(#loc46))
#loc3106 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/mul"(#loc46))
#loc3107 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add"(#loc47))
#loc3108 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3109 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/div"(#loc47))
#loc3110 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/mul"(#loc47))
#loc3111 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add"(#loc52))
#loc3112 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3113 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/mul"(#loc52))
#loc3114 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3115 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add"(#loc53))
#loc3116 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3117 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3118 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3119 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3120 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3121 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3122 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3123 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3124 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3125 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/sub"(#loc62))
#loc3126 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3127 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/mul"(#loc63))
#loc3128 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3129 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3130 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3131 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3132 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3133 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/mul"(#loc64))
#loc3134 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/mul"(#loc9))
#loc3135 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/add"(#loc12))
#loc3136 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3137 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3138 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3139 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/sub"(#loc62))
#loc3140 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3141 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/mul"(#loc36))
#loc3142 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3143 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3144 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3145 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/max"(#loc69))
#loc3146 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/div"(#loc5))
#loc3147 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3148 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/sub"(#loc70))
#loc3149 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/mul"(#loc70))
#loc3150 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/sqrt"(#loc71))
#loc3151 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/div"(#loc71))
#loc3152 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/abs"(#loc68))
#loc3153 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ge"(#loc68))
#loc3154 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3155 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3156 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/eq"(#loc68))
#loc3157 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3158 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3159 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3160 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3161 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3162 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3163 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/div"(#loc45))
#loc3164 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/sub"(#loc46))
#loc3165 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/mul"(#loc46))
#loc3166 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add"(#loc47))
#loc3167 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/rsqrt"(#loc47))
#loc3168 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/div"(#loc47))
#loc3169 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/mul"(#loc47))
#loc3170 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add"(#loc52))
#loc3171 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3172 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/mul"(#loc52))
#loc3173 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3174 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add"(#loc53))
#loc3175 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3176 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3177 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3178 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3179 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3180 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3181 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3182 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3183 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3184 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3185 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3186 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3187 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3188 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3189 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3190 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3191 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3192 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3193 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3194 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3195 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3196 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3197 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3198 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3199 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3200 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3201 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3202 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3203 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3204 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3205 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3206 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3207 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3208 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/add"(#loc12))
#loc3209 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3210 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3211 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3212 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3213 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3214 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/div"(#loc45))
#loc3215 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/sub"(#loc46))
#loc3216 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/mul"(#loc46))
#loc3217 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add"(#loc47))
#loc3218 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3219 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/div"(#loc47))
#loc3220 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/mul"(#loc47))
#loc3221 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add"(#loc52))
#loc3222 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3223 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/mul"(#loc52))
#loc3224 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3225 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add"(#loc53))
#loc3226 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3227 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3228 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3229 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3230 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3231 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3232 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3233 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3234 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3235 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/sub"(#loc62))
#loc3236 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3237 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/mul"(#loc63))
#loc3238 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3239 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3240 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3241 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3242 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3243 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/mul"(#loc64))
#loc3244 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/mul"(#loc9))
#loc3245 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/add"(#loc12))
#loc3246 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3247 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3248 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3249 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/sub"(#loc62))
#loc3250 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3251 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/mul"(#loc36))
#loc3252 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3253 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3254 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3255 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/max"(#loc69))
#loc3256 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/div"(#loc5))
#loc3257 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3258 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/sub"(#loc70))
#loc3259 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/mul"(#loc70))
#loc3260 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/sqrt"(#loc71))
#loc3261 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/div"(#loc71))
#loc3262 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/abs"(#loc68))
#loc3263 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ge"(#loc68))
#loc3264 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3265 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3266 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/eq"(#loc68))
#loc3267 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3268 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3269 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3270 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3271 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3272 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3273 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/div"(#loc45))
#loc3274 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/sub"(#loc46))
#loc3275 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/mul"(#loc46))
#loc3276 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add"(#loc47))
#loc3277 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/rsqrt"(#loc47))
#loc3278 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/div"(#loc47))
#loc3279 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/mul"(#loc47))
#loc3280 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add"(#loc52))
#loc3281 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3282 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/mul"(#loc52))
#loc3283 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3284 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add"(#loc53))
#loc3285 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3286 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3287 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3288 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3289 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3290 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3291 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3292 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3293 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3294 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3295 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3296 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3297 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3298 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3299 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3300 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3301 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3302 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3303 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3304 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3305 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3306 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3307 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3308 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3309 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3310 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3311 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3312 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3313 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3314 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3315 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3316 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3317 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3318 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/add"(#loc12))
#loc3319 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3320 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3321 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3322 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3323 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3324 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/div"(#loc45))
#loc3325 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/sub"(#loc46))
#loc3326 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/mul"(#loc46))
#loc3327 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add"(#loc47))
#loc3328 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3329 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/div"(#loc47))
#loc3330 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/mul"(#loc47))
#loc3331 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add"(#loc52))
#loc3332 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3333 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/mul"(#loc52))
#loc3334 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3335 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add"(#loc53))
#loc3336 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3337 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3338 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3339 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3340 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3341 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3342 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3343 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3344 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3345 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/sub"(#loc62))
#loc3346 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3347 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/mul"(#loc63))
#loc3348 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3349 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3350 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3351 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3352 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3353 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/mul"(#loc64))
#loc3354 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/mul"(#loc9))
#loc3355 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/add"(#loc12))
#loc3356 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3357 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3358 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3359 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/sub"(#loc62))
#loc3360 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3361 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/mul"(#loc36))
#loc3362 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3363 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3364 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3365 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/max"(#loc69))
#loc3366 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/div"(#loc5))
#loc3367 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3368 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/sub"(#loc70))
#loc3369 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/mul"(#loc70))
#loc3370 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/sqrt"(#loc71))
#loc3371 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/div"(#loc71))
#loc3372 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/abs"(#loc68))
#loc3373 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ge"(#loc68))
#loc3374 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3375 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3376 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/eq"(#loc68))
#loc3377 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3378 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3379 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3380 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3381 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3382 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3383 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/div"(#loc45))
#loc3384 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/sub"(#loc46))
#loc3385 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/mul"(#loc46))
#loc3386 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add"(#loc47))
#loc3387 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/rsqrt"(#loc47))
#loc3388 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/div"(#loc47))
#loc3389 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/mul"(#loc47))
#loc3390 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add"(#loc52))
#loc3391 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3392 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/mul"(#loc52))
#loc3393 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3394 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add"(#loc53))
#loc3395 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3396 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3397 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3398 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3399 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3400 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3401 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3402 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3403 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3404 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3405 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3406 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3407 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3408 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3409 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3410 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3411 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3412 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3413 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3414 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3415 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3416 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3417 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3418 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3419 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3420 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3421 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3422 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3423 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3424 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3425 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3426 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3427 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3428 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/add"(#loc12))
#loc3429 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3430 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3431 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3432 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3433 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3434 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/div"(#loc45))
#loc3435 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/sub"(#loc46))
#loc3436 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/mul"(#loc46))
#loc3437 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add"(#loc47))
#loc3438 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3439 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/div"(#loc47))
#loc3440 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/mul"(#loc47))
#loc3441 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add"(#loc52))
#loc3442 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3443 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/mul"(#loc52))
#loc3444 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3445 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add"(#loc53))
#loc3446 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3447 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3448 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3449 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3450 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3451 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3452 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3453 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3454 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3455 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/sub"(#loc62))
#loc3456 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3457 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/mul"(#loc63))
#loc3458 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3459 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3460 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3461 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3462 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3463 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/mul"(#loc64))
#loc3464 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/mul"(#loc9))
#loc3465 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/add"(#loc12))
#loc3466 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3467 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3468 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3469 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/sub"(#loc62))
#loc3470 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3471 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/mul"(#loc36))
#loc3472 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3473 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3474 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3475 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/max"(#loc69))
#loc3476 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/div"(#loc5))
#loc3477 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3478 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/sub"(#loc70))
#loc3479 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/mul"(#loc70))
#loc3480 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/sqrt"(#loc71))
#loc3481 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/div"(#loc71))
#loc3482 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/abs"(#loc68))
#loc3483 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ge"(#loc68))
#loc3484 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3485 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3486 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/eq"(#loc68))
#loc3487 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3488 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3489 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3490 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3491 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3492 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3493 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/div"(#loc45))
#loc3494 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/sub"(#loc46))
#loc3495 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/mul"(#loc46))
#loc3496 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add"(#loc47))
#loc3497 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/rsqrt"(#loc47))
#loc3498 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/div"(#loc47))
#loc3499 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/mul"(#loc47))
#loc3500 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add"(#loc52))
#loc3501 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3502 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/mul"(#loc52))
#loc3503 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3504 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add"(#loc53))
#loc3505 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3506 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3507 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3508 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3509 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3510 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3511 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3512 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3513 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3514 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3515 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3516 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3517 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3518 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3519 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3520 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3521 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3522 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3523 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3524 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3525 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3526 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3527 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3528 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3529 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3530 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3531 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3532 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3533 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3534 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3535 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3536 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3537 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3538 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/add"(#loc12))
#loc3539 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3540 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3541 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3542 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3543 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3544 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/div"(#loc45))
#loc3545 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/sub"(#loc46))
#loc3546 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/mul"(#loc46))
#loc3547 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add"(#loc47))
#loc3548 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3549 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/div"(#loc47))
#loc3550 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/mul"(#loc47))
#loc3551 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add"(#loc52))
#loc3552 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3553 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/mul"(#loc52))
#loc3554 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3555 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add"(#loc53))
#loc3556 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3557 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3558 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3559 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3560 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3561 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3562 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3563 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3564 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3565 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/sub"(#loc62))
#loc3566 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3567 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/mul"(#loc63))
#loc3568 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3569 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3570 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3571 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3572 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3573 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/mul"(#loc64))
#loc3574 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/mul"(#loc9))
#loc3575 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/add"(#loc12))
#loc3576 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3577 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3578 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3579 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/sub"(#loc62))
#loc3580 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3581 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/mul"(#loc36))
#loc3582 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3583 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3584 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3585 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/max"(#loc69))
#loc3586 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/div"(#loc5))
#loc3587 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3588 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/sub"(#loc70))
#loc3589 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/mul"(#loc70))
#loc3590 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/sqrt"(#loc71))
#loc3591 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/div"(#loc71))
#loc3592 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/abs"(#loc68))
#loc3593 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ge"(#loc68))
#loc3594 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3595 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3596 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/eq"(#loc68))
#loc3597 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3598 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3599 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3600 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3601 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3602 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3603 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/div"(#loc45))
#loc3604 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/sub"(#loc46))
#loc3605 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/mul"(#loc46))
#loc3606 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add"(#loc47))
#loc3607 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/rsqrt"(#loc47))
#loc3608 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/div"(#loc47))
#loc3609 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/mul"(#loc47))
#loc3610 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add"(#loc52))
#loc3611 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3612 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/mul"(#loc52))
#loc3613 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3614 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add"(#loc53))
#loc3615 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3616 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3617 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3618 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3619 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3620 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3621 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3622 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3623 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3624 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3625 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3626 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3627 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3628 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3629 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3630 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3631 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3632 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3633 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3634 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3635 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3636 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3637 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3638 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3639 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3640 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3641 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3642 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3643 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3644 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3645 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3646 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3647 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3648 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/add"(#loc12))
#loc3649 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3650 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3651 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3652 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3653 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3654 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/div"(#loc45))
#loc3655 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/sub"(#loc46))
#loc3656 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/mul"(#loc46))
#loc3657 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add"(#loc47))
#loc3658 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3659 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/div"(#loc47))
#loc3660 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/mul"(#loc47))
#loc3661 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add"(#loc52))
#loc3662 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3663 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/mul"(#loc52))
#loc3664 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3665 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add"(#loc53))
#loc3666 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3667 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3668 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3669 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3670 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3671 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3672 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3673 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3674 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3675 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/sub"(#loc62))
#loc3676 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3677 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/mul"(#loc63))
#loc3678 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3679 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3680 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3681 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/add"(#loc53))
#loc3682 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3683 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/mul"(#loc64))
#loc3684 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/mul"(#loc9))
#loc3685 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/add"(#loc12))
#loc3686 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc67))
#loc3687 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=False]"(#loc66))
#loc3688 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3689 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/sub"(#loc62))
#loc3690 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=False]"(#loc36))
#loc3691 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/mul"(#loc36))
#loc3692 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_sum[axes=(1, 2, 3)]"(#loc36))
#loc3693 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 2048, 2048) broadcast_dimensions=()]"(#loc35))
#loc3694 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(1, 2, 3)]"(#loc36))
#loc3695 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/max"(#loc69))
#loc3696 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/div"(#loc5))
#loc3697 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc70))
#loc3698 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/sub"(#loc70))
#loc3699 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/mul"(#loc70))
#loc3700 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/sqrt"(#loc71))
#loc3701 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/div"(#loc71))
#loc3702 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/abs"(#loc68))
#loc3703 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ge"(#loc68))
#loc3704 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_max[axes=(1, 2, 3)]"(#loc68))
#loc3705 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc68))
#loc3706 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/eq"(#loc68))
#loc3707 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=float32 weak_type=False]"(#loc68))
#loc3708 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_sum[axes=(1, 2, 3)]"(#loc68))
#loc3709 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3710 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3711 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3712 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3713 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/div"(#loc45))
#loc3714 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/sub"(#loc46))
#loc3715 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/mul"(#loc46))
#loc3716 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add"(#loc47))
#loc3717 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/rsqrt"(#loc47))
#loc3718 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/div"(#loc47))
#loc3719 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/mul"(#loc47))
#loc3720 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add"(#loc52))
#loc3721 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3722 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/mul"(#loc52))
#loc3723 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3724 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add"(#loc53))
#loc3725 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3726 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/slice[start_indices=(0, 0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3727 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/squeeze[dimensions=(1,)]"(#loc54))
#loc3728 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/slice[start_indices=(0, 1, 0, 0, 0, 0) limit_indices=(1, 2, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3729 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/slice[start_indices=(0, 2, 0, 0, 0, 0) limit_indices=(1, 3, 1, 2048, 32, 64) strides=(1, 1, 1, 1, 1, 1)]"(#loc54))
#loc3730 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._shard_blnh/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3731 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False,) name=softplus keep_unused=False inline=False]"(#loc37))
#loc3732 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc37))
#loc3733 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 1, 1, 1, 64) broadcast_dimensions=(0, 4)]"(#loc55))
#loc3734 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/mul"(#loc55))
#loc3735 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3736 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/div"(#loc56))
#loc3737 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/tanh"(#loc56))
#loc3738 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/sub"(#loc56))
#loc3739 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/mul"(#loc56))
#loc3740 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc57))
#loc3741 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/mul"(#loc38))
#loc3742 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc38))
#loc3743 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/ge"(#loc38))
#loc3744 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc3745 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/reduce_max[axes=(4,)]"(#loc58))
#loc3746 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc3747 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/sub"(#loc58))
#loc3748 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/exp"(#loc58))
#loc3749 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/reduce_sum[axes=(4,)]"(#loc58))
#loc3750 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/div"(#loc58))
#loc3751 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/mul"(#loc58))
#loc3752 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/integer_pow[y=-2]"(#loc58))
#loc3753 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 1) broadcast_dimensions=()]"(#loc58))
#loc3754 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc58))
#loc3755 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3756 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3757 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._shard_bld/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3758 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/add"(#loc12))
#loc3759 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc61))
#loc3760 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/convert_element_type[new_dtype=float32 weak_type=False]"(#loc45))
#loc3761 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc45))
#loc3762 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 1) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc3763 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc45))
#loc3764 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/div"(#loc45))
#loc3765 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/sub"(#loc46))
#loc3766 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/mul"(#loc46))
#loc3767 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add"(#loc47))
#loc3768 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/rsqrt"(#loc47))
#loc3769 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/div"(#loc47))
#loc3770 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/mul"(#loc47))
#loc3771 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add"(#loc52))
#loc3772 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc52))
#loc3773 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/mul"(#loc52))
#loc3774 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 1, 2048) broadcast_dimensions=(0, 3)]"(#loc53))
#loc3775 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add"(#loc53))
#loc3776 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc3777 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc3778 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 1, 1, 8192) broadcast_dimensions=(0, 3)]"(#loc59))
#loc3779 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/add"(#loc59))
#loc3780 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/mul"(#loc60))
#loc3781 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/add"(#loc60))
#loc3782 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/tanh"(#loc60))
#loc3783 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/sub"(#loc60))
#loc3784 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/convert_element_type[new_dtype=float32 weak_type=True]"(#loc62))
#loc3785 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/sub"(#loc62))
#loc3786 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc63))
#loc3787 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/mul"(#loc63))
#loc3788 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc64))
#loc3789 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc39))
#loc3790 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/mul"(#loc71))
#loc3791 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3792 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/neg"(#loc70))
#loc3793 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3794 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3795 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3796 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3797 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/add_any"(#loc39))
#loc3798 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/add_any"(#loc36))
#loc3799 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/div"(#loc68))
#loc3800 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3801 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/mul"(#loc68))
#loc3802 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3803 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/select_n"(#loc68))
#loc3804 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/neg"(#loc68))
#loc3805 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/add_any"(#loc68))
#loc3806 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3807 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/mul"(#loc71))
#loc3808 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3809 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/neg"(#loc70))
#loc3810 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3811 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3812 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3813 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3814 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/add_any"(#loc36))
#loc3815 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/div"(#loc68))
#loc3816 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3817 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/mul"(#loc68))
#loc3818 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3819 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/select_n"(#loc68))
#loc3820 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/neg"(#loc68))
#loc3821 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/add_any"(#loc68))
#loc3822 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3823 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/mul"(#loc71))
#loc3824 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3825 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/neg"(#loc70))
#loc3826 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3827 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3828 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3829 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3830 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/add_any"(#loc36))
#loc3831 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/div"(#loc68))
#loc3832 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3833 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/mul"(#loc68))
#loc3834 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3835 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/select_n"(#loc68))
#loc3836 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/neg"(#loc68))
#loc3837 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/add_any"(#loc68))
#loc3838 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3839 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/mul"(#loc71))
#loc3840 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3841 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/neg"(#loc70))
#loc3842 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3843 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3844 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3845 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3846 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/add_any"(#loc36))
#loc3847 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/div"(#loc68))
#loc3848 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3849 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/mul"(#loc68))
#loc3850 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3851 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/select_n"(#loc68))
#loc3852 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/neg"(#loc68))
#loc3853 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/add_any"(#loc68))
#loc3854 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3855 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/mul"(#loc71))
#loc3856 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3857 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/neg"(#loc70))
#loc3858 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3859 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3860 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3861 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3862 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/add_any"(#loc36))
#loc3863 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/div"(#loc68))
#loc3864 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3865 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/mul"(#loc68))
#loc3866 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3867 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/select_n"(#loc68))
#loc3868 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/neg"(#loc68))
#loc3869 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/add_any"(#loc68))
#loc3870 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3871 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/mul"(#loc71))
#loc3872 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3873 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/neg"(#loc70))
#loc3874 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3875 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3876 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3877 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3878 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/add_any"(#loc36))
#loc3879 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/div"(#loc68))
#loc3880 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3881 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/mul"(#loc68))
#loc3882 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3883 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/select_n"(#loc68))
#loc3884 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/neg"(#loc68))
#loc3885 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/add_any"(#loc68))
#loc3886 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3887 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/mul"(#loc71))
#loc3888 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3889 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/neg"(#loc70))
#loc3890 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3891 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3892 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3893 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3894 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/add_any"(#loc36))
#loc3895 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/div"(#loc68))
#loc3896 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3897 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/mul"(#loc68))
#loc3898 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3899 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/select_n"(#loc68))
#loc3900 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/neg"(#loc68))
#loc3901 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/add_any"(#loc68))
#loc3902 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3903 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/mul"(#loc71))
#loc3904 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3905 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/neg"(#loc70))
#loc3906 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3907 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3908 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3909 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3910 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/add_any"(#loc36))
#loc3911 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/div"(#loc68))
#loc3912 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3913 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/mul"(#loc68))
#loc3914 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3915 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/select_n"(#loc68))
#loc3916 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/neg"(#loc68))
#loc3917 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/add_any"(#loc68))
#loc3918 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3919 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/mul"(#loc71))
#loc3920 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3921 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/neg"(#loc70))
#loc3922 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3923 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3924 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3925 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3926 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/add_any"(#loc36))
#loc3927 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/div"(#loc68))
#loc3928 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3929 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/mul"(#loc68))
#loc3930 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3931 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/select_n"(#loc68))
#loc3932 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/neg"(#loc68))
#loc3933 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/add_any"(#loc68))
#loc3934 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3935 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/mul"(#loc71))
#loc3936 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3937 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/neg"(#loc70))
#loc3938 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3939 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3940 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3941 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3942 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/add_any"(#loc36))
#loc3943 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/div"(#loc68))
#loc3944 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3945 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/mul"(#loc68))
#loc3946 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3947 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/select_n"(#loc68))
#loc3948 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/neg"(#loc68))
#loc3949 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/add_any"(#loc68))
#loc3950 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3951 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/mul"(#loc71))
#loc3952 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3953 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/neg"(#loc70))
#loc3954 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3955 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3956 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3957 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3958 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/add_any"(#loc36))
#loc3959 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/div"(#loc68))
#loc3960 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3961 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/mul"(#loc68))
#loc3962 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3963 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/select_n"(#loc68))
#loc3964 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/neg"(#loc68))
#loc3965 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/add_any"(#loc68))
#loc3966 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3967 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/mul"(#loc71))
#loc3968 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3969 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/neg"(#loc70))
#loc3970 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3971 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3972 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3973 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3974 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/add_any"(#loc36))
#loc3975 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/div"(#loc68))
#loc3976 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3977 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/mul"(#loc68))
#loc3978 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3979 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/select_n"(#loc68))
#loc3980 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/neg"(#loc68))
#loc3981 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/add_any"(#loc68))
#loc3982 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3983 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/mul"(#loc71))
#loc3984 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc3985 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/neg"(#loc70))
#loc3986 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_sum[axes=(2, 3)]"(#loc70))
#loc3987 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc3988 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc3989 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc3990 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/add_any"(#loc36))
#loc3991 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/div"(#loc68))
#loc3992 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc3993 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/mul"(#loc68))
#loc3994 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc3995 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/select_n"(#loc68))
#loc3996 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/neg"(#loc68))
#loc3997 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/add_any"(#loc68))
#loc3998 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc3999 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/mul"(#loc71))
#loc4000 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc4001 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/neg"(#loc70))
#loc4002 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_sum[axes=(2, 3)]"(#loc70))
#loc4003 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc4004 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc4005 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc4006 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/add_any"(#loc36))
#loc4007 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/div"(#loc68))
#loc4008 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc4009 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/mul"(#loc68))
#loc4010 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc4011 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/select_n"(#loc68))
#loc4012 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/neg"(#loc68))
#loc4013 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/add_any"(#loc68))
#loc4014 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc4015 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/mul"(#loc71))
#loc4016 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc4017 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/neg"(#loc70))
#loc4018 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_sum[axes=(2, 3)]"(#loc70))
#loc4019 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc4020 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc4021 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc4022 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/add_any"(#loc36))
#loc4023 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/div"(#loc68))
#loc4024 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc4025 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/mul"(#loc68))
#loc4026 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc4027 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/select_n"(#loc68))
#loc4028 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/neg"(#loc68))
#loc4029 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/add_any"(#loc68))
#loc4030 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc4031 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/mul"(#loc71))
#loc4032 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc36))
#loc4033 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/neg"(#loc70))
#loc4034 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_sum[axes=(2, 3)]"(#loc70))
#loc4035 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reshape[new_sizes=(1, 1, 1, 1) dimensions=None]"(#loc70))
#loc4036 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/reduce_sum[axes=(0, 1, 2, 3)]"(#loc70))
#loc4037 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc70))
#loc4038 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/add_any"(#loc36))
#loc4039 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/div"(#loc68))
#loc4040 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0,)]"(#loc68))
#loc4041 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/mul"(#loc68))
#loc4042 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc68))
#loc4043 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/select_n"(#loc68))
#loc4044 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/neg"(#loc68))
#loc4045 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/add_any"(#loc68))
#loc4046 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc66))
#loc4047 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/pad[padding_config=((0, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc40))
#loc4048 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_pad keep_unused=False inline=False]"(#loc40))
#loc4049 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/add_any"(#loc40))
#loc4050 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/mul"(#loc9))
#loc4051 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/mul"(#loc64))
#loc4052 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4053 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4054 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4055 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4056 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc51))
#loc4057 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc4058 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4059 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4060 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4061 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4062 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4063 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4064 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4065 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4066 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4067 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4068 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4069 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4070 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4071 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4072 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4073 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4074 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4075 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4076 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4077 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4078 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/neg"(#loc46))
#loc4079 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4080 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4081 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add_any"(#loc46))
#loc4082 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/layer_norm/add_any"(#loc45))
#loc4083 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/add_any"(#loc12))
#loc4084 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4085 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4086 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4087 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4088 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4089 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4090 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4091 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4092 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4093 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4094 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4095 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4096 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4097 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4098 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4099 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4100 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/add_any"(#loc54))
#loc4101 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4102 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4103 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4104 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4105 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4106 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4107 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4108 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4109 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4110 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4111 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4112 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4113 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4114 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4115 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4116 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4117 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4118 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4119 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4120 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4121 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/neg"(#loc46))
#loc4122 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4123 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4124 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add_any"(#loc46))
#loc4125 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/layer_norm/add_any"(#loc45))
#loc4126 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4127 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4128 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4129 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4130 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4131 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4132 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4133 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4134 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4135 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4136 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4137 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4138 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4139 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4140 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4141 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4142 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4143 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4144 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4145 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4146 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4147 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4148 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4149 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4150 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/neg"(#loc46))
#loc4151 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4152 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4153 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add_any"(#loc46))
#loc4154 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/layer_norm/add_any"(#loc45))
#loc4155 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/add_any"(#loc12))
#loc4156 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4157 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4158 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4159 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4160 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4161 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4162 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4163 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4164 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4165 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4166 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4167 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4168 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4169 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4170 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4171 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4172 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/add_any"(#loc54))
#loc4173 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4174 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4175 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4176 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4177 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4178 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4179 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4180 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4181 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4182 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4183 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4184 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4185 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4186 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4187 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4188 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4189 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4190 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4191 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4192 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4193 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/neg"(#loc46))
#loc4194 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4195 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4196 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add_any"(#loc46))
#loc4197 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/layer_norm/add_any"(#loc45))
#loc4198 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4199 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4200 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4201 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4202 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4203 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4204 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4205 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4206 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4207 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4208 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4209 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4210 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4211 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4212 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4213 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4214 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4215 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4216 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4217 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4218 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4219 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4220 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4221 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4222 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/neg"(#loc46))
#loc4223 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4224 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4225 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add_any"(#loc46))
#loc4226 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/layer_norm/add_any"(#loc45))
#loc4227 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/add_any"(#loc12))
#loc4228 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4229 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4230 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4231 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4232 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4233 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4234 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4235 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4236 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4237 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4238 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4239 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4240 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4241 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4242 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4243 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4244 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/add_any"(#loc54))
#loc4245 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4246 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4247 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4248 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4249 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4250 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4251 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4252 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4253 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4254 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4255 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4256 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4257 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4258 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4259 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4260 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4261 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4262 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4263 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4264 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4265 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/neg"(#loc46))
#loc4266 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4267 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4268 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add_any"(#loc46))
#loc4269 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/layer_norm/add_any"(#loc45))
#loc4270 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4271 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4272 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4273 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4274 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4275 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4276 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4277 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4278 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4279 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4280 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4281 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4282 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4283 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4284 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4285 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4286 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4287 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4288 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4289 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4290 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4291 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4292 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4293 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4294 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/neg"(#loc46))
#loc4295 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4296 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4297 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add_any"(#loc46))
#loc4298 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/layer_norm/add_any"(#loc45))
#loc4299 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/add_any"(#loc12))
#loc4300 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4301 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4302 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4303 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4304 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4305 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4306 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4307 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4308 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4309 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4310 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4311 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4312 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4313 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4314 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4315 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4316 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/add_any"(#loc54))
#loc4317 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4318 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4319 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4320 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4321 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4322 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4323 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4324 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4325 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4326 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4327 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4328 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4329 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4330 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4331 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4332 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4333 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4334 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4335 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4336 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4337 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/neg"(#loc46))
#loc4338 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4339 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4340 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add_any"(#loc46))
#loc4341 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/layer_norm/add_any"(#loc45))
#loc4342 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4343 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4344 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4345 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4346 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4347 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4348 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4349 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4350 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4351 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4352 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4353 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4354 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4355 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4356 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4357 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4358 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4359 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4360 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4361 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4362 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4363 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4364 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4365 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4366 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/neg"(#loc46))
#loc4367 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4368 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4369 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add_any"(#loc46))
#loc4370 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/layer_norm/add_any"(#loc45))
#loc4371 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/add_any"(#loc12))
#loc4372 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4373 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4374 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4375 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4376 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4377 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4378 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4379 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4380 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4381 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4382 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4383 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4384 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4385 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4386 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4387 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4388 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/add_any"(#loc54))
#loc4389 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4390 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4391 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4392 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4393 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4394 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4395 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4396 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4397 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4398 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4399 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4400 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4401 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4402 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4403 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4404 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4405 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4406 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4407 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4408 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4409 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/neg"(#loc46))
#loc4410 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4411 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4412 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add_any"(#loc46))
#loc4413 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/layer_norm/add_any"(#loc45))
#loc4414 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4415 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4416 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4417 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4418 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4419 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4420 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4421 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4422 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4423 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4424 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4425 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4426 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4427 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4428 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4429 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4430 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4431 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4432 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4433 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4434 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4435 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4436 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4437 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4438 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/neg"(#loc46))
#loc4439 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4440 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4441 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add_any"(#loc46))
#loc4442 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/layer_norm/add_any"(#loc45))
#loc4443 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/add_any"(#loc12))
#loc4444 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4445 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4446 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4447 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4448 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4449 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4450 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4451 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4452 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4453 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4454 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4455 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4456 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4457 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4458 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4459 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4460 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/add_any"(#loc54))
#loc4461 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4462 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4463 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4464 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4465 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4466 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4467 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4468 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4469 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4470 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4471 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4472 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4473 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4474 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4475 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4476 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4477 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4478 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4479 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4480 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4481 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/neg"(#loc46))
#loc4482 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4483 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4484 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add_any"(#loc46))
#loc4485 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/layer_norm/add_any"(#loc45))
#loc4486 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4487 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4488 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4489 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4490 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4491 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4492 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4493 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4494 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4495 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4496 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4497 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4498 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4499 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4500 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4501 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4502 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4503 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4504 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4505 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4506 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4507 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4508 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4509 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4510 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/neg"(#loc46))
#loc4511 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4512 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4513 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add_any"(#loc46))
#loc4514 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/layer_norm/add_any"(#loc45))
#loc4515 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/add_any"(#loc12))
#loc4516 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4517 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4518 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4519 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4520 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4521 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4522 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4523 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4524 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4525 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4526 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4527 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4528 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4529 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4530 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4531 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4532 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/add_any"(#loc54))
#loc4533 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4534 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4535 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4536 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4537 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4538 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4539 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4540 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4541 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4542 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4543 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4544 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4545 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4546 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4547 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4548 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4549 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4550 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4551 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4552 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4553 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/neg"(#loc46))
#loc4554 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4555 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4556 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add_any"(#loc46))
#loc4557 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/layer_norm/add_any"(#loc45))
#loc4558 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4559 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4560 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4561 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4562 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4563 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4564 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4565 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4566 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4567 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4568 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4569 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4570 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4571 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4572 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4573 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4574 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4575 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4576 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4577 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4578 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4579 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4580 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4581 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4582 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/neg"(#loc46))
#loc4583 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4584 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4585 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add_any"(#loc46))
#loc4586 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/layer_norm/add_any"(#loc45))
#loc4587 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/add_any"(#loc12))
#loc4588 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4589 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4590 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4591 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4592 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4593 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4594 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4595 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4596 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4597 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4598 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4599 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4600 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4601 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4602 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4603 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4604 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/add_any"(#loc54))
#loc4605 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4606 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4607 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4608 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4609 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4610 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4611 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4612 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4613 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4614 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4615 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4616 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4617 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4618 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4619 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4620 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4621 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4622 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4623 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4624 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4625 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/neg"(#loc46))
#loc4626 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4627 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4628 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add_any"(#loc46))
#loc4629 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/layer_norm/add_any"(#loc45))
#loc4630 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4631 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4632 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4633 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4634 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4635 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4636 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4637 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4638 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4639 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4640 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4641 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4642 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4643 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4644 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4645 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4646 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4647 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4648 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4649 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4650 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4651 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4652 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4653 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4654 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/neg"(#loc46))
#loc4655 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4656 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4657 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add_any"(#loc46))
#loc4658 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/layer_norm/add_any"(#loc45))
#loc4659 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/add_any"(#loc12))
#loc4660 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4661 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4662 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4663 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4664 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4665 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4666 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4667 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4668 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4669 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4670 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4671 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4672 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4673 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4674 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4675 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4676 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/add_any"(#loc54))
#loc4677 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4678 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4679 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4680 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4681 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4682 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4683 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4684 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4685 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4686 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4687 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4688 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4689 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4690 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4691 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4692 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4693 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4694 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4695 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4696 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4697 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/neg"(#loc46))
#loc4698 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4699 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4700 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add_any"(#loc46))
#loc4701 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/layer_norm/add_any"(#loc45))
#loc4702 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4703 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4704 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4705 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4706 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4707 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4708 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4709 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4710 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4711 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4712 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4713 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4714 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4715 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4716 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4717 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4718 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4719 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4720 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4721 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4722 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4723 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4724 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4725 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4726 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/neg"(#loc46))
#loc4727 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4728 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4729 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add_any"(#loc46))
#loc4730 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/layer_norm/add_any"(#loc45))
#loc4731 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/add_any"(#loc12))
#loc4732 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4733 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4734 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4735 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4736 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4737 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4738 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4739 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4740 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4741 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4742 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4743 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4744 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4745 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4746 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4747 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4748 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/add_any"(#loc54))
#loc4749 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4750 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4751 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4752 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4753 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4754 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4755 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4756 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4757 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4758 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4759 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4760 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4761 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4762 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4763 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4764 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4765 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4766 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4767 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4768 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4769 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/neg"(#loc46))
#loc4770 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4771 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4772 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add_any"(#loc46))
#loc4773 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/layer_norm/add_any"(#loc45))
#loc4774 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4775 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4776 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4777 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4778 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4779 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4780 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4781 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4782 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4783 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4784 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4785 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4786 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4787 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4788 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4789 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4790 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4791 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4792 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4793 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4794 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4795 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4796 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4797 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4798 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/neg"(#loc46))
#loc4799 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4800 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4801 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add_any"(#loc46))
#loc4802 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/layer_norm/add_any"(#loc45))
#loc4803 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/add_any"(#loc12))
#loc4804 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4805 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4806 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4807 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4808 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4809 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4810 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4811 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4812 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4813 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4814 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4815 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4816 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4817 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4818 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4819 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4820 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/add_any"(#loc54))
#loc4821 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4822 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4823 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4824 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4825 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4826 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4827 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4828 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4829 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4830 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4831 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4832 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4833 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4834 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4835 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4836 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4837 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4838 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4839 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4840 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4841 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/neg"(#loc46))
#loc4842 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4843 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4844 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add_any"(#loc46))
#loc4845 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/layer_norm/add_any"(#loc45))
#loc4846 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4847 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4848 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4849 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4850 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4851 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4852 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4853 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4854 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4855 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4856 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4857 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4858 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4859 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4860 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4861 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4862 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4863 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4864 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4865 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4866 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4867 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4868 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4869 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4870 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/neg"(#loc46))
#loc4871 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4872 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4873 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add_any"(#loc46))
#loc4874 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/layer_norm/add_any"(#loc45))
#loc4875 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/add_any"(#loc12))
#loc4876 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4877 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4878 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4879 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4880 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4881 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4882 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4883 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4884 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4885 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4886 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4887 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4888 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4889 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4890 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4891 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4892 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/add_any"(#loc54))
#loc4893 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4894 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4895 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4896 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4897 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4898 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4899 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4900 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4901 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4902 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4903 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4904 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4905 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4906 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4907 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4908 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4909 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4910 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4911 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4912 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4913 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/neg"(#loc46))
#loc4914 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4915 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4916 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add_any"(#loc46))
#loc4917 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/layer_norm/add_any"(#loc45))
#loc4918 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4919 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4920 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4921 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4922 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4923 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4924 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4925 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4926 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4927 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc4928 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4929 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4930 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4931 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4932 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4933 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4934 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4935 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4936 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4937 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4938 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4939 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4940 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4941 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4942 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/neg"(#loc46))
#loc4943 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4944 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4945 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add_any"(#loc46))
#loc4946 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/layer_norm/add_any"(#loc45))
#loc4947 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/add_any"(#loc12))
#loc4948 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4949 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4950 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc4951 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4952 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc4953 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc4954 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc4955 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc4956 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc4957 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc4958 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc4959 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc4960 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc4961 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc4962 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4963 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4964 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/add_any"(#loc54))
#loc4965 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc4966 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc4967 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc4968 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc4969 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc4970 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc4971 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4972 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc4973 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4974 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4975 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4976 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc4977 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc4978 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc4979 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc4980 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc4981 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc4982 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc4983 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc4984 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc4985 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/neg"(#loc46))
#loc4986 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc4987 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc4988 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add_any"(#loc46))
#loc4989 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/layer_norm/add_any"(#loc45))
#loc4990 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc4991 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc4992 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc4993 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc4994 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc4995 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc4996 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc4997 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc4998 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc4999 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc5000 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5001 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc5002 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5003 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5004 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5005 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc5006 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc5007 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc5008 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc5009 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc5010 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc5011 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc5012 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc5013 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc5014 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/neg"(#loc46))
#loc5015 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc5016 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc5017 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add_any"(#loc46))
#loc5018 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/layer_norm/add_any"(#loc45))
#loc5019 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/add_any"(#loc12))
#loc5020 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5021 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5022 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc5023 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5024 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc5025 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc5026 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc5027 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc5028 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc5029 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc5030 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc5031 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc5032 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc5033 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc5034 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5035 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5036 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/add_any"(#loc54))
#loc5037 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc5038 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc5039 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc5040 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc5041 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc5042 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5043 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5044 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc5045 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5046 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5047 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5048 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc5049 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc5050 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc5051 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc5052 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc5053 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc5054 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc5055 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc5056 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc5057 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/neg"(#loc46))
#loc5058 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc5059 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc5060 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add_any"(#loc46))
#loc5061 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/layer_norm/add_any"(#loc45))
#loc5062 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc5063 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5064 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5065 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5066 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5067 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc5068 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc5069 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc5070 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc5071 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc5072 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5073 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc5074 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5075 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5076 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5077 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc5078 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc5079 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc5080 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc5081 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc5082 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc5083 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc5084 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc5085 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc5086 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/neg"(#loc46))
#loc5087 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc5088 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc5089 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add_any"(#loc46))
#loc5090 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/layer_norm/add_any"(#loc45))
#loc5091 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/add_any"(#loc12))
#loc5092 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5093 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5094 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc5095 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5096 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc5097 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc5098 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc5099 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc5100 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc5101 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc5102 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc5103 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc5104 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc5105 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc5106 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5107 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5108 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/add_any"(#loc54))
#loc5109 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc5110 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc5111 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc5112 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc5113 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc5114 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5115 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5116 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc5117 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5118 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5119 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5120 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc5121 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc5122 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc5123 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc5124 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc5125 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc5126 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc5127 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc5128 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc5129 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/neg"(#loc46))
#loc5130 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc5131 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc5132 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add_any"(#loc46))
#loc5133 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/layer_norm/add_any"(#loc45))
#loc5134 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/reduce_sum[axes=(2,)]"(#loc53))
#loc5135 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5136 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5137 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/bias/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5138 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5139 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/activation/add_any"(#loc60))
#loc5140 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/reduce_sum[axes=(2,)]"(#loc59))
#loc5141 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/reshape[new_sizes=(1, 1, 1, 8192) dimensions=None]"(#loc59))
#loc5142 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/reduce_sum[axes=(0, 1, 2)]"(#loc59))
#loc5143 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/bias/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc59))
#loc5144 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5145 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc5146 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5147 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5148 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5149 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc5150 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc5151 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc5152 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc5153 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc5154 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc5155 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc5156 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc5157 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc5158 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/neg"(#loc46))
#loc5159 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc5160 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc5161 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add_any"(#loc46))
#loc5162 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/layer_norm/add_any"(#loc45))
#loc5163 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/add_any"(#loc12))
#loc5164 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5165 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5166 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/broadcast_in_dim[shape=(1, 1, 1, 2048, 32, 64) broadcast_dimensions=(0, 2, 3, 4, 5)]"(#loc54))
#loc5167 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (2, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5168 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=float32 weak_type=False]"(#loc58))
#loc5169 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/reshape[new_sizes=(1, 1, 32, 2048, 1) dimensions=None]"(#loc58))
#loc5170 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/neg"(#loc58))
#loc5171 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/reduce_sum[axes=(0, 1, 4)]"(#loc58))
#loc5172 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048) broadcast_dimensions=(2, 3)]"(#loc58))
#loc5173 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 2, 3)]"(#loc58))
#loc5174 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/add_any"(#loc58))
#loc5175 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc38))
#loc5176 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc57))
#loc5177 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._cap_logits/add_any"(#loc56))
#loc5178 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5179 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (1, 1, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5180 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/add_any"(#loc54))
#loc5181 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(2, 3)]"(#loc55))
#loc5182 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reshape[new_sizes=(1, 1, 1, 1, 64) dimensions=None]"(#loc55))
#loc5183 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/reduce_sum[axes=(0, 1, 2, 3)]"(#loc55))
#loc5184 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc55))
#loc5185 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=softplus keep_unused=False inline=False]"(#loc37))
#loc5186 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/pad[padding_config=((0, 0, 0), (0, 2, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc54))
#loc5187 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5188 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(2,)]"(#loc53))
#loc5189 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc53))
#loc5190 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc53))
#loc5191 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc53))
#loc5192 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(2,)]"(#loc52))
#loc5193 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reshape[new_sizes=(1, 1, 1, 2048) dimensions=None]"(#loc52))
#loc5194 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(0, 1, 2)]"(#loc52))
#loc5195 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc52))
#loc5196 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(3,)]"(#loc47))
#loc5197 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc47))
#loc5198 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(0, 1, 3)]"(#loc45))
#loc5199 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 2048) broadcast_dimensions=(2,)]"(#loc45))
#loc5200 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(0, 1, 2)]"(#loc45))
#loc5201 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/neg"(#loc46))
#loc5202 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reduce_sum[axes=(3,)]"(#loc46))
#loc5203 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/reshape[new_sizes=(1, 1, 2048, 1) dimensions=None]"(#loc46))
#loc5204 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add_any"(#loc46))
#loc5205 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/layer_norm/add_any"(#loc45))
#loc5206 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue, UnspecifiedValue) resource_env=None donated_invars=(False, False) name=_where keep_unused=False inline=False]"(#loc44))
#loc5207 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/broadcast_in_dim[shape=(1, 1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc31))
#loc5208 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1, 2, 3), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]"(#loc31))
#loc5209 = loc("/add"(#loc31))
#loc5210 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/add_any"(#loc41))
#loc5211 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=_pad keep_unused=False inline=False]"(#loc27))
#loc5212 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reduce_sum[axes=(0, 1, 2)]"(#loc26))
#loc5213 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=(2, 3)]"(#loc26))
#loc5214 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 1, 1, 2048, 2048) dimensions=None]"(#loc25))
#loc5215 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/transpose[permutation=(0, 2, 1, 3, 4)]"(#loc24))
#loc5216 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/reshape[new_sizes=(1, 2048, 2048) dimensions=None]"(#loc23))
#loc5217 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc5218 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5219 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/pad[padding_config=((0, 0, 0), (0, 0, 0))]"(#loc10))
#loc5220 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/sharding_constraint[sharding=GSPMDSharding({maximal device=0}) resource_env=ResourceEnv(Mesh(device_ids=array([[[[0]]]]), axis_names=('stage', 'replica', 'data', 'mdl')), ()) unconstrained_dims=set()]"(#loc1))
#loc5221 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/mul"(#loc9))
#loc5222 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/einsum/...y,yz->...z/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=_einsum keep_unused=False inline=False]"(#loc8))
#loc5223 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/transpose[permutation=(1, 0)]"(#loc6))
#loc5224 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/add_any"(#loc6))
#loc5225 = loc("pjit(_wrapped_step_fn)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc100))
#loc5226 = loc("pjit(_wrapped_step_fn)/jit(main)/sub"(#loc101))
#loc5227 = loc("pjit(_wrapped_step_fn)/jit(main)/div"(#loc101))
#loc5228 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc38))
#loc5229 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc102))
#loc5230 = loc("pjit(_wrapped_step_fn)/jit(main)/lt"(#loc103))
#loc5231 = loc("pjit(_wrapped_step_fn)/jit(main)/ge"(#loc103))
#loc5232 = loc("pjit(_wrapped_step_fn)/jit(main)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc103))
#loc5233 = loc("pjit(_wrapped_step_fn)/jit(main)/lt"(#loc104))
#loc5234 = loc("pjit(_wrapped_step_fn)/jit(main)/sub"(#loc104))
#loc5235 = loc("pjit(_wrapped_step_fn)/jit(main)/cos"(#loc105))
#loc5236 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc106))
#loc5237 = loc("pjit(_wrapped_step_fn)/jit(main)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc104))
#loc5238 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc107))
#loc5239 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0,)]"(#loc107))
#loc5240 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0, 1)]"(#loc107))
#loc5241 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0, 1, 2)]"(#loc107))
#loc5242 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc108))
#loc5243 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0, 1)]"(#loc4))
#loc5244 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0, 1, 2, 3, 4)]"(#loc107))
#loc5245 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0, 1, 2, 3)]"(#loc107))
#loc5246 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc109))
#loc5247 = loc("pjit(_wrapped_step_fn)/jit(main)/concatenate[dimension=0]"(#loc109))
#loc5248 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0,)]"(#loc109))
#loc5249 = loc("pjit(_wrapped_step_fn)/jit(main)/sqrt"(#loc71))
#loc5250 = loc("pjit(_wrapped_step_fn)/jit(main)/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=isfinite keep_unused=False inline=False]"(#loc110))
#loc5251 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_and[axes=()]"(#loc110))
#loc5252 = loc("pjit(_wrapped_step_fn)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc111))
#loc5253 = loc("pjit(_wrapped_step_fn)/jit(main)/div"(#loc5))
#loc5254 = loc("pjit(_wrapped_step_fn)/jit(main)/min"(#loc112))
#loc5255 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc113))
#loc5256 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc65))
#loc5257 = loc("pjit(_wrapped_step_fn)/jit(main)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False, False) name=nan_to_num keep_unused=False inline=False]"(#loc114))
#loc5258 = loc("pjit(_wrapped_step_fn)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc115))
#loc5259 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc92))
#loc5260 = loc("pjit(_wrapped_step_fn)/jit(main)/sub"(#loc116))
#loc5261 = loc("pjit(_wrapped_step_fn)/jit(main)/pow"(#loc116))
#loc5262 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc117))
#loc5263 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc117))
#loc5264 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc118))
#loc5265 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc119))
#loc5266 = loc("pjit(_wrapped_step_fn)/jit(main)/sqrt"(#loc119))
#loc5267 = loc("pjit(_wrapped_step_fn)/jit(main)/div"(#loc119))
#loc5268 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc120))
#loc5269 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0,)]"(#loc121))
#loc5270 = loc("pjit(_wrapped_step_fn)/jit(main)/div"(#loc121))
#loc5271 = loc("pjit(_wrapped_step_fn)/jit(main)/max"(#loc69))
#loc5272 = loc("pjit(_wrapped_step_fn)/jit(main)/div"(#loc122))
#loc5273 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(0, 1)]"(#loc121))
#loc5274 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(1,)]"(#loc123))
#loc5275 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc124))
#loc5276 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc42))
#loc5277 = loc("pjit(_wrapped_step_fn)/jit(main)/lt"(#loc31))
#loc5278 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=(0,)]"(#loc117))
#loc5279 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1, 1) broadcast_dimensions=(0,)]"(#loc117))
#loc5280 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc117))
#loc5281 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc117))
#loc5282 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(1,)]"(#loc121))
#loc5283 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1) broadcast_dimensions=(0,)]"(#loc122))
#loc5284 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(1, 2)]"(#loc121))
#loc5285 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1, 1) broadcast_dimensions=(0,)]"(#loc122))
#loc5286 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(1, 2, 3, 4)]"(#loc121))
#loc5287 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc122))
#loc5288 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_sum[axes=(1, 2, 3)]"(#loc121))
#loc5289 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 1, 1, 1) broadcast_dimensions=(0,)]"(#loc122))
#loc5290 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(2048,) broadcast_dimensions=()]"(#loc125))
#loc5291 = loc("pjit(_wrapped_step_fn)/jit(main)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc125))
#loc5292 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(2048, 2048) broadcast_dimensions=()]"(#loc125))
#loc5293 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(51200,) broadcast_dimensions=()]"(#loc125))
#loc5294 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(2048, 51200) broadcast_dimensions=()]"(#loc125))
#loc5295 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=()]"(#loc125))
#loc5296 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=()]"(#loc125))
#loc5297 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=()]"(#loc125))
#loc5298 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=()]"(#loc125))
#loc5299 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=()]"(#loc125))
#loc5300 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=()]"(#loc125))
#loc5301 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=()]"(#loc125))
#loc5302 = loc("pjit(_wrapped_step_fn)/jit(main)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc126))
#loc5303 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc127))
#loc5304 = loc("pjit(_wrapped_step_fn)/jit(main)/concatenate[dimension=0]"(#loc127))
#loc5305 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc128))
#loc5306 = loc("pjit(_wrapped_step_fn)/jit(main)/add"(#loc129))
#loc5307 = loc("pjit(_wrapped_step_fn)/jit(main)/abs"(#loc130))
#loc5308 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_max[axes=(0,)]"(#loc130))
#loc5309 = loc("pjit(_wrapped_step_fn)/jit(main)/max"(#loc130))
#loc5310 = loc("pjit(_wrapped_step_fn)/jit(main)/mul"(#loc131))
#loc5311 = loc("pjit(_wrapped_step_fn)/jit(main)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc131))
#loc5312 = loc("pjit(_wrapped_step_fn)/jit(main)/concatenate[dimension=0]"(#loc131))
#loc5313 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_max[axes=(0, 1)]"(#loc130))
#loc5314 = loc("pjit(_wrapped_step_fn)/jit(main)/div"(#loc132))
#loc5315 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_max[axes=(0, 1, 2)]"(#loc130))
#loc5316 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_max[axes=(0, 1, 2, 3, 4)]"(#loc130))
#loc5317 = loc("pjit(_wrapped_step_fn)/jit(main)/reduce_max[axes=(0, 1, 2, 3)]"(#loc130))
#loc5318 = loc("pjit(_wrapped_step_fn)/jit(main)/dynamic_slice[slice_sizes=(1,)]"(#loc133))
#loc5319 = loc("pjit(_wrapped_step_fn)/jit(main)/squeeze[dimensions=(0,)]"(#loc133))
#loc5320 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/jit(_one_hot)/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc7))
#loc5321 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/jit(_one_hot)/iota[dtype=int32 shape=(1, 1, 51200) dimension=2]"(#loc7))
#loc5322 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/jit(_one_hot)/eq"(#loc7))
#loc5323 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/jit(_one_hot)/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc7))
#loc5324 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc5325 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/jit(_one_hot)/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc11))
#loc5326 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/jit(_one_hot)/iota[dtype=int32 shape=(1, 1, 2048) dimension=2]"(#loc11))
#loc5327 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/jit(_one_hot)/eq"(#loc11))
#loc5328 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/jit(_one_hot)/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc11))
#loc5329 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc5330 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/broadcast_in_dim[shape=(5, 2) broadcast_dimensions=()]"(#loc27))
#loc5331 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc27))
#loc5332 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc27))
#loc5333 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/concatenate[dimension=0]"(#loc27))
#loc5334 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc27))
#loc5335 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/pad[padding_config=((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc27))
#loc5336 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/jit(remainder)/convert_element_type[new_dtype=int32 weak_type=False]"(#loc30))
#loc5337 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/jit(remainder)/eq"(#loc30))
#loc5338 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/jit(remainder)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc30))
#loc5339 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/jit(remainder)/lt"(#loc30))
#loc5340 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/jit(remainder)/jit(_where)/select_n"(#loc30))
#loc5341 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/convert_element_type[new_dtype=int32 weak_type=False]"(#loc31))
#loc5342 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/eq"(#loc30))
#loc5343 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc31))
#loc5344 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/lt"(#loc31))
#loc5345 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/jit(_where)/select_n"(#loc30))
#loc5346 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5347 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5348 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5349 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5350 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5351 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5352 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5353 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5354 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5355 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5356 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5357 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5358 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5359 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5360 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5361 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5362 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5363 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5364 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5365 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5366 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5367 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5368 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5369 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5370 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5371 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5372 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5373 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5374 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5375 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5376 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5377 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5378 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc39))
#loc5379 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/broadcast_in_dim[shape=(4, 2) broadcast_dimensions=()]"(#loc40))
#loc5380 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/convert_element_type[new_dtype=bfloat16 weak_type=False]"(#loc40))
#loc5381 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc40))
#loc5382 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/concatenate[dimension=0]"(#loc40))
#loc5383 = loc("pjit(_wrapped_step_fn)/jit(main)/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]"(#loc40))
#loc5384 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(remainder)/rem"(#loc30))
#loc5385 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(remainder)/lt"(#loc30))
#loc5386 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(remainder)/ne"(#loc30))
#loc5387 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(remainder)/and"(#loc30))
#loc5388 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(remainder)/add"(#loc30))
#loc5389 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(remainder)/select_n"(#loc30))
#loc5390 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/rem"(#loc31))
#loc5391 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/lt"(#loc31))
#loc5392 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/ne"(#loc31))
#loc5393 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/and"(#loc31))
#loc5394 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/add"(#loc31))
#loc5395 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/select_n"(#loc31))
#loc5396 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/jit(_where)/select_n"(#loc44))
#loc5397 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(0,)]"(#loc51))
#loc5398 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/select_n"(#loc51))
#loc5399 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=(0,)]"(#loc51))
#loc5400 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5401 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5402 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(0,)]"(#loc51))
#loc5403 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5404 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5405 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5406 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5407 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5408 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5409 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5410 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5411 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5412 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5413 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5414 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5415 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5416 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5417 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5418 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5419 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=(0,)]"(#loc51))
#loc5420 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5421 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=(0,)]"(#loc51))
#loc5422 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5423 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(0,)]"(#loc51))
#loc5424 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=(0,)]"(#loc51))
#loc5425 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5426 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5427 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5428 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5429 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5430 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5431 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5432 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5433 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5434 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5435 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5436 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5437 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5438 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5439 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5440 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5441 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5442 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5443 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5444 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5445 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5446 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5447 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5448 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5449 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5450 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5451 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5452 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5453 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5454 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5455 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5456 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5457 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5458 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5459 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5460 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5461 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5462 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5463 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5464 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5465 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5466 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5467 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5468 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5469 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5470 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5471 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5472 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5473 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5474 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5475 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5476 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5477 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5478 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5479 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5480 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5481 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5482 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5483 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5484 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5485 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5486 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5487 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5488 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5489 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5490 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5491 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5492 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5493 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5494 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5495 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5496 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5497 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5498 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5499 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5500 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5501 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5502 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5503 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5504 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5505 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5506 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5507 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5508 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5509 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5510 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5511 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5512 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5513 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5514 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5515 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5516 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5517 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5518 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5519 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5520 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5521 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5522 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5523 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5524 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5525 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5526 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5527 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5528 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5529 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5530 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5531 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5532 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5533 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5534 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5535 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5536 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5537 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5538 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5539 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5540 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5541 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5542 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5543 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5544 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5545 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5546 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5547 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5548 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5549 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5550 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5551 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5552 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5553 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5554 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5555 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5556 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5557 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5558 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5559 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5560 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5561 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5562 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5563 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5564 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5565 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5566 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5567 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5568 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5569 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5570 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5571 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5572 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5573 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5574 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5575 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5576 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5577 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5578 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5579 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5580 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5581 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5582 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5583 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5584 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5585 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5586 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5587 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5588 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5589 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5590 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5591 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5592 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5593 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5594 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5595 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5596 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5597 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5598 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5599 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5600 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5601 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5602 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5603 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5604 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5605 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5606 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5607 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5608 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5609 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5610 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5611 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5612 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5613 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5614 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5615 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5616 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5617 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5618 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5619 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5620 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5621 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5622 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5623 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5624 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5625 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5626 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5627 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5628 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5629 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5630 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5631 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5632 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5633 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5634 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5635 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5636 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5637 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5638 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5639 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5640 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5641 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5642 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5643 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5644 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5645 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5646 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5647 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5648 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5649 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5650 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5651 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5652 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5653 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5654 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5655 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5656 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5657 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5658 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5659 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5660 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5661 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5662 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5663 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5664 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5665 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5666 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5667 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5668 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5669 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5670 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5671 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5672 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5673 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5674 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5675 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5676 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5677 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5678 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5679 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5680 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5681 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5682 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5683 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5684 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5685 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5686 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5687 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5688 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5689 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5690 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5691 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5692 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5693 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5694 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5695 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5696 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5697 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5698 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5699 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5700 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5701 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5702 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5703 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5704 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5705 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5706 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5707 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5708 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5709 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5710 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5711 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5712 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5713 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5714 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5715 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5716 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5717 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5718 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5719 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5720 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5721 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5722 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5723 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5724 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5725 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5726 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5727 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5728 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5729 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5730 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5731 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5732 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5733 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5734 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5735 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5736 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5737 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5738 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5739 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5740 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5741 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/pad[padding_config=((1, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc40))
#loc5742 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/pad[padding_config=((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc40))
#loc5743 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/select_n"(#loc39))
#loc5744 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc5745 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/reduce_max[axes=(2,)]"(#loc83))
#loc5746 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc83))
#loc5747 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/sub"(#loc83))
#loc5748 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/exp"(#loc83))
#loc5749 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/reduce_sum[axes=(2,)]"(#loc83))
#loc5750 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/log"(#loc83))
#loc5751 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(_one_hot)/broadcast_in_dim[shape=(1, 2048, 1) broadcast_dimensions=(0, 1)]"(#loc85))
#loc5752 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(_one_hot)/iota[dtype=int32 shape=(1, 1, 51200) dimension=2]"(#loc85))
#loc5753 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(_one_hot)/eq"(#loc7))
#loc5754 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(_one_hot)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc85))
#loc5755 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/iota[dtype=int32 shape=(1, 2048, 51200) dimension=2]"(#loc88))
#loc5756 = loc("pjit(_wrapped_step_fn)/jit(main)/jvp(xformer_lm.apply)/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/reduce[computation=<function _compute_argminmax.<locals>.reducer_fn at 0x7fccb48925f0> consts=() dimensions=(2,)]"(#loc88))
#loc5757 = loc("/gt"(#loc88))
#loc5758 = loc("/ne"(#loc88))
#loc5759 = loc("/or"(#loc88))
#loc5760 = loc("/eq"(#loc88))
#loc5761 = loc("/lt"(#loc88))
#loc5762 = loc("/and"(#loc88))
#loc5763 = loc("/select_n"(#loc88))
#loc5764 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/neg"(#loc83))
#loc5765 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/reduce_sum[axes=(2,)]"(#loc83))
#loc5766 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/reshape[new_sizes=(1, 2048, 1) dimensions=None]"(#loc83))
#loc5767 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/div"(#loc83))
#loc5768 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/reduce_sum[axes=(0, 2)]"(#loc83))
#loc5769 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc83))
#loc5770 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/broadcast_in_dim[shape=(1, 2048, 51200) broadcast_dimensions=(0, 1)]"(#loc83))
#loc5771 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/mul"(#loc83))
#loc5772 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/jit(log_softmax)/add_any"(#loc83))
#loc5773 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((0, 1), (0, 1)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc5774 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(1, 0)]"(#loc8))
#loc5775 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm.compute_loss/softmax/softmax.get_logits/logits_ffn/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((2,), (1,)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc5776 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/convert_element_type[new_dtype=int32 weak_type=False]"(#loc30))
#loc5777 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/eq"(#loc30))
#loc5778 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc30))
#loc5779 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/rem"(#loc30))
#loc5780 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/ne"(#loc30))
#loc5781 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/lt"(#loc30))
#loc5782 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/and"(#loc30))
#loc5783 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/add"(#loc30))
#loc5784 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/select_n"(#loc30))
#loc5785 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(remainder)/jit(_where)/select_n"(#loc30))
#loc5786 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/convert_element_type[new_dtype=int32 weak_type=False]"(#loc31))
#loc5787 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/eq"(#loc30))
#loc5788 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc31))
#loc5789 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/rem"(#loc31))
#loc5790 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/ne"(#loc31))
#loc5791 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/lt"(#loc31))
#loc5792 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/and"(#loc31))
#loc5793 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/add"(#loc31))
#loc5794 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/select_n"(#loc31))
#loc5795 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_iteration_inputs/jit(remainder)/jit(_where)/select_n"(#loc30))
#loc5796 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(_where)/select_n"(#loc44))
#loc5797 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(0,)]"(#loc51))
#loc5798 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/select_n"(#loc51))
#loc5799 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=(0,)]"(#loc51))
#loc5800 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(0,)]"(#loc51))
#loc5801 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=(0,)]"(#loc51))
#loc5802 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=(0,)]"(#loc51))
#loc5803 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(0,)]"(#loc51))
#loc5804 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=(0,)]"(#loc51))
#loc5805 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5806 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5807 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5808 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5809 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5810 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5811 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5812 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5813 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5814 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5815 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5816 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5817 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5818 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5819 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5820 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5821 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5822 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5823 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5824 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5825 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5826 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5827 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5828 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5829 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5830 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5831 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5832 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5833 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5834 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5835 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5836 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5837 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5838 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5839 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5840 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5841 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5842 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5843 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5844 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5845 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5846 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5847 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5848 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5849 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5850 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5851 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5852 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5853 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5854 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5855 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5856 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5857 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5858 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5859 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5860 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5861 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5862 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5863 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5864 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5865 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5866 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5867 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5868 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5869 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5870 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5871 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5872 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5873 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5874 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5875 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5876 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5877 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5878 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5879 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5880 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5881 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5882 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5883 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5884 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5885 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5886 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5887 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5888 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5889 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5890 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5891 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5892 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5893 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5894 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5895 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5896 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5897 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5898 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5899 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5900 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5901 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5902 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5903 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5904 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5905 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5906 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5907 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5908 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5909 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5910 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5911 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5912 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5913 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5914 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5915 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5916 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5917 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5918 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5919 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5920 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5921 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5922 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5923 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5924 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5925 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5926 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5927 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5928 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5929 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5930 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5931 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5932 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5933 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5934 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5935 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5936 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5937 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5938 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5939 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5940 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5941 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5942 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5943 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5944 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5945 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5946 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5947 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5948 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5949 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5950 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5951 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5952 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5953 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5954 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5955 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5956 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5957 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5958 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5959 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5960 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5961 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5962 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5963 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5964 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5965 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5966 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5967 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5968 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5969 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5970 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc5971 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc5972 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc5973 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc5974 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc5975 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc5976 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5977 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5978 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5979 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc5980 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc5981 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc5982 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc5983 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc5984 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc5985 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc5986 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc5987 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc5988 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc5989 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc5990 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc5991 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc5992 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc5993 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc5994 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc5995 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc5996 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc5997 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc5998 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc5999 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6000 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6001 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6002 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6003 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6004 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6005 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6006 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6007 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6008 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6009 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6010 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6011 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6012 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6013 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6014 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6015 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6016 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6017 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6018 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6019 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6020 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6021 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6022 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6023 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6024 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6025 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6026 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6027 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6028 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6029 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6030 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6031 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6032 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6033 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6034 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6035 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6036 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6037 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6038 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6039 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6040 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6041 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6042 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6043 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6044 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6045 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6046 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6047 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6048 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6049 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6050 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6051 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6052 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6053 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6054 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6055 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6056 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6057 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6058 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6059 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6060 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6061 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6062 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6063 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6064 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6065 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6066 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6067 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6068 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6069 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6070 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6071 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6072 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6073 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6074 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6075 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6076 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6077 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6078 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6079 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6080 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6081 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6082 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6083 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6084 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6085 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6086 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6087 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6088 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6089 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6090 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6091 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6092 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6093 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6094 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6095 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6096 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6097 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6098 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6099 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6100 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6101 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6102 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6103 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6104 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6105 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6106 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6107 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6108 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6109 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6110 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6111 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6112 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6113 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6114 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6115 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6116 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6117 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6118 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6119 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6120 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6121 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6122 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6123 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6124 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6125 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6126 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6127 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6128 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6129 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6130 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6131 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6132 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6133 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6134 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6135 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6136 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6137 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6138 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6139 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6140 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6141 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6142 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6143 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6144 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6145 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6146 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6147 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6148 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6149 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6150 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6151 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6152 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6153 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6154 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6155 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6156 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6157 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6158 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6159 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6160 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6161 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6162 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6163 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6164 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6165 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6166 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6167 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6168 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6169 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6170 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6171 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6172 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6173 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6174 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6175 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6176 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6177 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6178 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6179 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6180 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6181 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6182 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6183 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6184 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6185 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6186 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6187 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6188 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6189 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6190 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6191 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6192 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6193 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6194 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6195 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6196 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6197 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6198 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6199 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6200 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6201 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6202 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6203 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6204 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6205 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6206 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6207 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6208 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6209 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6210 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6211 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6212 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6213 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6214 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6215 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6216 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6217 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6218 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6219 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6220 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6221 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6222 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6223 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6224 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6225 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6226 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6227 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6228 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6229 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6230 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6231 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6232 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6233 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6234 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6235 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6236 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6237 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6238 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6239 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6240 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (3,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6241 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 5, 2, 3)]"(#loc8))
#loc6242 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False) name=logaddexp keep_unused=False inline=False]"(#loc37))
#loc6243 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/eq"(#loc37))
#loc6244 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc37))
#loc6245 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc37))
#loc6246 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/select_n"(#loc37))
#loc6247 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/sub"(#loc37))
#loc6248 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/exp"(#loc37))
#loc6249 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/max"(#loc37))
#loc6250 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/sub"(#loc37))
#loc6251 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/ne"(#loc37))
#loc6252 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/add"(#loc37))
#loc6253 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/abs"(#loc37))
#loc6254 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/neg"(#loc37))
#loc6255 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/exp"(#loc37))
#loc6256 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/log1p"(#loc37))
#loc6257 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/jit(logaddexp)/select_n"(#loc37))
#loc6258 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (4,)), ((0, 1, 3), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6259 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/squeeze[dimensions=(2,)]"(#loc38))
#loc6260 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(0, 1, 3, 4)]"(#loc38))
#loc6261 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6262 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=(1, 2, 3, 4)]"(#loc38))
#loc6263 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/select_n"(#loc38))
#loc6264 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((2,), (4,)), ((0, 1, 3), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6265 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6266 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((4, 3), (3, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6267 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6268 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc39))
#loc6269 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/eq"(#loc39))
#loc6270 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/select_n"(#loc39))
#loc6271 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/pad[padding_config=((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc40))
#loc6272 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/slice[start_indices=(0, 0, 0, 0) limit_indices=(2, 1, 2048, 2048) strides=(1, 1, 1, 1)]"(#loc40))
#loc6273 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/pad[padding_config=((-1, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc40))
#loc6274 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline._get_new_loop_state/jit(_pad)/slice[start_indices=(0, 0, 0, 0) limit_indices=(1, 1, 2048, 2048) strides=(1, 1, 1, 1)]"(#loc40))
#loc6275 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=()]"(#loc51))
#loc6276 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/eq"(#loc51))
#loc6277 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6278 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6279 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6280 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=()]"(#loc51))
#loc6281 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=()]"(#loc51))
#loc6282 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6283 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6284 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6285 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=()]"(#loc51))
#loc6286 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6287 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6288 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=()]"(#loc51))
#loc6289 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6290 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6291 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6292 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6293 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6294 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6295 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6296 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6297 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6298 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=()]"(#loc51))
#loc6299 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6300 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6301 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_15/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6302 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/jit(_where)/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=()]"(#loc51))
#loc6303 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6304 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6305 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6306 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6307 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6308 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6309 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6310 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6311 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6312 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6313 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6314 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6315 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6316 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6317 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6318 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6319 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6320 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6321 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6322 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_14/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6323 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6324 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6325 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6326 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6327 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6328 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6329 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6330 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6331 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6332 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6333 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6334 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6335 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6336 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6337 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6338 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6339 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6340 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6341 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6342 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_13/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6343 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6344 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6345 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6346 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6347 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6348 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6349 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6350 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6351 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6352 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6353 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6354 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6355 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6356 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6357 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6358 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6359 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6360 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6361 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6362 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_12/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6363 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6364 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6365 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6366 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6367 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6368 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6369 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6370 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6371 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6372 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6373 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6374 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6375 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6376 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6377 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6378 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6379 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6380 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6381 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6382 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_11/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6383 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6384 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6385 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6386 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6387 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6388 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6389 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6390 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6391 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6392 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6393 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6394 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6395 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6396 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6397 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6398 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6399 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6400 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6401 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6402 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_10/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6403 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6404 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6405 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6406 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6407 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6408 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6409 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6410 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6411 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6412 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6413 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6414 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6415 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6416 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6417 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6418 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6419 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6420 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6421 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6422 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_9/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6423 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6424 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6425 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6426 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6427 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6428 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6429 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6430 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6431 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6432 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6433 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6434 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6435 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6436 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6437 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6438 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6439 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6440 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6441 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6442 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_8/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6443 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6444 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6445 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6446 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6447 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6448 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6449 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6450 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6451 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6452 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6453 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6454 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6455 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6456 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6457 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6458 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6459 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6460 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6461 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6462 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_7/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6463 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6464 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6465 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6466 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6467 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6468 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6469 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6470 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6471 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6472 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6473 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6474 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6475 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6476 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6477 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6478 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6479 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6480 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6481 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6482 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_6/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6483 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6484 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6485 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6486 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6487 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6488 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6489 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6490 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6491 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6492 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6493 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6494 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6495 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6496 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6497 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6498 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6499 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6500 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6501 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6502 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_5/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6503 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6504 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6505 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6506 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6507 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6508 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6509 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6510 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6511 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6512 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6513 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6514 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6515 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6516 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6517 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6518 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6519 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6520 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6521 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6522 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_4/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6523 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6524 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6525 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6526 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6527 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6528 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6529 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6530 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6531 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6532 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6533 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6534 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6535 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6536 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6537 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6538 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6539 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6540 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6541 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6542 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_3/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6543 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6544 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6545 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6546 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6547 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6548 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6549 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6550 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6551 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6552 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6553 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6554 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6555 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6556 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6557 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6558 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6559 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6560 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6561 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6562 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_2/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6563 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6564 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6565 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6566 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6567 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6568 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6569 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6570 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6571 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6572 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6573 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6574 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6575 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6576 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6577 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6578 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6579 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6580 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6581 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6582 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_1/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6583 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6584 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6585 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer2/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6586 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6587 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(0, 2, 1)]"(#loc8))
#loc6588 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/ff_layer/ffn_layer1/linear/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6589 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((1, 2), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6590 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/post/einsum/ABNH,DNH->ABD/jit(_einsum)/dot_general[dimension_numbers=(((3,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6591 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/transpose[permutation=(0, 1, 3, 4, 2)]"(#loc8))
#loc6592 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((3,), (4,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6593 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/pv_einsum/BNTS,BSNH->BTNH/jit(_einsum)/dot_general[dimension_numbers=(((4,), (3,)), ((0, 1, 2), (0, 1, 2))) precision=None preferred_element_type=None]"(#loc8))
#loc6594 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/broadcast_in_dim[shape=(1, 1, 32, 2048, 2048) broadcast_dimensions=()]"(#loc38))
#loc6595 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/jit(_where)/eq"(#loc38))
#loc6596 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((3,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6597 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/transpose[permutation=(0, 1, 3, 2, 4)]"(#loc8))
#loc6598 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._atten_logits/qk_einsum/BTNH,BSNH->BNTS/jit(_einsum)/dot_general[dimension_numbers=(((4,), (2,)), ((0, 1, 2), (0, 1, 3))) precision=None preferred_element_type=None]"(#loc8))
#loc6599 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/self_attention._dot_atten/self_attention._scale_query/per_dim_scale/jit(softplus)/mul"(#loc37))
#loc6600 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((1, 2, 3), (1, 3, 4)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6601 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/dot_general[dimension_numbers=(((4, 5), (1, 2)), ((0,), (0,))) precision=None preferred_element_type=None]"(#loc8))
#loc6602 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/pipeline.body_fprop/vmap(body.body_fn)/body.layer_fprop/body/x_layers_0/self_attention/combined_qkv/einsum/ABD,KDNH->KABNH/jit(_einsum)/transpose[permutation=(0, 1, 4, 2, 3)]"(#loc8))
#loc6603 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(_where)/broadcast_in_dim[shape=(1, 1, 2048, 2048) broadcast_dimensions=()]"(#loc44))
#loc6604 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/while/body/remat/pipeline._scan_fn/jit(_where)/eq"(#loc44))
#loc6605 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/pad[padding_config=((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0))]"(#loc27))
#loc6606 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/transformer/pipeline/jit(_pad)/slice[start_indices=(0, 0, 0, 0, 0) limit_indices=(1, 1, 1, 2048, 2048) strides=(1, 1, 1, 1, 1)]"(#loc27))
#loc6607 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((0, 1), (0, 1)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc6608 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/position_emb/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(1, 0)]"(#loc8))
#loc6609 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/einsum/...y,yz->...z/jit(_einsum)/dot_general[dimension_numbers=(((0, 1), (0, 1)), ((), ())) precision=None preferred_element_type=None]"(#loc8))
#loc6610 = loc("pjit(_wrapped_step_fn)/jit(main)/transpose(jvp(xformer_lm.apply))/xformer_lm/xformer_lm.compute_predictions/lm/lm._prepare_input/softmax.emb_lookup/einsum/...y,yz->...z/jit(_einsum)/transpose[permutation=(1, 0)]"(#loc8))
#loc6611 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc103))
#loc6612 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/select_n"(#loc103))
#loc6613 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/select_n"(#loc39))
#loc6614 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(isfinite)/is_finite"(#loc110))
#loc6615 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=isnan keep_unused=False inline=False]"(#loc114))
#loc6616 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/convert_element_type[new_dtype=float32 weak_type=False]"(#loc114))
#loc6617 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc114))
#loc6618 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/eq"(#loc114))
#loc6619 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(isnan)/ne"(#loc114))
#loc6620 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(_where)/broadcast_in_dim[shape=(2048,) broadcast_dimensions=()]"(#loc114))
#loc6621 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(_where)/select_n"(#loc114))
#loc6622 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(_where)/broadcast_in_dim[shape=(2048, 2048) broadcast_dimensions=()]"(#loc114))
#loc6623 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(_where)/broadcast_in_dim[shape=(51200,) broadcast_dimensions=()]"(#loc114))
#loc6624 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(_where)/broadcast_in_dim[shape=(2048, 51200) broadcast_dimensions=()]"(#loc114))
#loc6625 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(nan_to_num)/jit(_where)/select_n"(#loc39))
#loc6626 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=isnan keep_unused=False inline=False]"(#loc114))
#loc6627 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/convert_element_type[new_dtype=float32 weak_type=False]"(#loc114))
#loc6628 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/pjit[in_shardings=(UnspecifiedValue, UnspecifiedValue, UnspecifiedValue) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False, False, False) name=_where keep_unused=False inline=False]"(#loc114))
#loc6629 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/eq"(#loc114))
#loc6630 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(isnan)/ne"(#loc114))
#loc6631 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(8192,) broadcast_dimensions=()]"(#loc114))
#loc6632 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=(1,)]"(#loc114))
#loc6633 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/select_n"(#loc114))
#loc6634 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(2048, 8192) broadcast_dimensions=()]"(#loc114))
#loc6635 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=(1, 2)]"(#loc114))
#loc6636 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(2048,) broadcast_dimensions=()]"(#loc114))
#loc6637 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=(1,)]"(#loc114))
#loc6638 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(8192, 2048) broadcast_dimensions=()]"(#loc114))
#loc6639 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=(1, 2)]"(#loc114))
#loc6640 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(3, 2048, 32, 64) broadcast_dimensions=()]"(#loc114))
#loc6641 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=(1, 2, 3, 4)]"(#loc114))
#loc6642 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(64,) broadcast_dimensions=()]"(#loc114))
#loc6643 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=(1,)]"(#loc114))
#loc6644 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(2048, 32, 64) broadcast_dimensions=()]"(#loc114))
#loc6645 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=(1, 2, 3)]"(#loc114))
#loc6646 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc39))
#loc6647 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(nan_to_num))/jit(_where)/select_n"(#loc39))
#loc6648 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(_where))/convert_element_type[new_dtype=float32 weak_type=False]"(#loc103))
#loc6649 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(_where))/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc103))
#loc6650 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(_where))/select_n"(#loc103))
#loc6651 = loc("pjit(_wrapped_step_fn)/jit(main)/vmap(jit(_where))/select_n"(#loc39))
#loc6652 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(2048,) broadcast_dimensions=()]"(#loc125))
#loc6653 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/select_n"(#loc125))
#loc6654 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(2048, 2048) broadcast_dimensions=()]"(#loc125))
#loc6655 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(51200,) broadcast_dimensions=()]"(#loc125))
#loc6656 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(2048, 51200) broadcast_dimensions=()]"(#loc125))
#loc6657 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 8192) broadcast_dimensions=()]"(#loc125))
#loc6658 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 8192) broadcast_dimensions=()]"(#loc125))
#loc6659 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 2048) broadcast_dimensions=()]"(#loc125))
#loc6660 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 8192, 2048) broadcast_dimensions=()]"(#loc125))
#loc6661 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 3, 2048, 32, 64) broadcast_dimensions=()]"(#loc125))
#loc6662 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 64) broadcast_dimensions=()]"(#loc125))
#loc6663 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1, 2048, 32, 64) broadcast_dimensions=()]"(#loc125))
#loc6664 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/select_n"(#loc30))
#loc6665 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/broadcast_in_dim[shape=(1,) broadcast_dimensions=()]"(#loc126))
#loc6666 = loc("pjit(_wrapped_step_fn)/jit(main)/jit(_where)/select_n"(#loc126))